<commit_before>package isolated

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("create-security-group command", func() {
	var (
		dir               string
		securityGroupName string
		tempPath          string
		userName          string
	)

	BeforeEach(func() {
		userName = helpers.LoginCF()
		securityGroupName = helpers.NewSecurityGroupName()
		dir = helpers.TempDirAbsolutePath("", "security-group")

		tempPath = filepath.Join(dir, "tmpfile.json")
		err := ioutil.WriteFile(tempPath, []byte(`[{
			"protocol": "tcp",
			"destination": "10.0.11.0/24",
			"ports": "80,443",
			"description": "Allow http and https traffic from ZoneA"
		}]`), 0666)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		err := os.RemoveAll(dir)
		Expect(err).NotTo(HaveOccurred())
	})

	Describe("help", func() {
		When("--help flag is set", func() {
			It("appears in cf help -a", func() {
				session := helpers.CF("help", "-a")
				Eventually(session).Should(Exit(0))
				Expect(session).To(HaveCommandInCategoryWithDescription("create-security-group", "SECURITY GROUP", "Create a security group"))
			})

			It("Displays command usage to output", func() {
				session := helpers.CF("create-security-group", "--help")
				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Say("create-security-group - Create a security group"))
				Eventually(session).Should(Say("USAGE:"))
				Eventually(session).Should(Say(`cf create-security-group SECURITY_GROUP PATH_TO_JSON_RULES_FILE`))
				Eventually(session).Should(Say(`The provided path can be an absolute or relative path to a file. The file should have`))
				Eventually(session).Should(Say(`a single array with JSON objects inside describing the rules. The JSON Base Object is`))
				Eventually(session).Should(Say(`omitted and only the square brackets and associated child object are required in the file.`))

				Eventually(session).Should(Say(`Valid json file example:`))
				Eventually(session).Should(Say(`\s+\[`))
				Eventually(session).Should(Say(`\s+{`))
				Eventually(session).Should(Say(`\s+"protocol": "tcp",`))
				Eventually(session).Should(Say(`\s+"destination": "10.0.11.0/24",`))
				Eventually(session).Should(Say(`\s+"ports": "80,443",`))
				Eventually(session).Should(Say(`\s+"description": "Allow http and https traffic from ZoneA"`))
				Eventually(session).Should(Say(`\s+}`))
				Eventually(session).Should(Say(`\s+\]`))
				Eventually(session).Should(Say("SEE ALSO:"))
				Eventually(session).Should(Say("bind-running-security-group, bind-security-group, bind-staging-security-group, security-groups"))
				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("not logged in", func() {
		It("fails with the appropriate errors", func() {
			helpers.CheckEnvironmentTargetedCorrectly(false, false, "", "create-security-group", "some-group", tempPath)
		})
	})

	When("the environment is set up correctly", func() {
		When("the security group name is not provided", func() {
			It("tells the user that the security group name is required, prints help text, and exits 1", func() {
				session := helpers.CF("create-security-group")

				Eventually(session.Err).Should(Say("Incorrect Usage: the required arguments `SECURITY_GROUP` and `PATH_TO_JSON_RULES_FILE` were not provided"))
				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("the security group path is not valid", func() {
			It("tells the user that the security group path does not exist, prints help text, and exits 1", func() {
				session := helpers.CF("create-security-group", securityGroupName, "/invalid/path")

				Eventually(session.Err).Should(Say("Incorrect Usage: The specified path '/invalid/path' does not exist."))
				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("the security group path has invalid JSON", func() {
			BeforeEach(func() {
				err := ioutil.WriteFile(tempPath, []byte("Invalid JSON!"), 0666)
				Expect(err).ToNot(HaveOccurred())
			})

			It("tells the user that the security group path is required, prints help text, and exits 1", func() {
				session := helpers.CF("create-security-group", securityGroupName, tempPath)
				Eventually(session).Should(Say("Creating security group %s as %s...", securityGroupName, userName))
				Eventually(session.Err).Should(Say("Incorrect json format: %s", strings.ReplaceAll(tempPath, "\\", "\\\\")))

				Eventually(session.Err).Should(Say("Valid json file example:"))
				Eventually(session.Err).Should(Say(`\[`))
				Eventually(session.Err).Should(Say(`\s+{`))
				Eventually(session.Err).Should(Say(`\s+"protocol": "tcp",`))
				Eventually(session.Err).Should(Say(`\s+"destination": "10.244.1.18",`))
				Eventually(session.Err).Should(Say(`\s+"ports": "3306"`))
				Eventually(session.Err).Should(Say(`\s+}`))
				Eventually(session.Err).Should(Say(`\]`))
				Eventually(session).Should(Exit(1))
			})
		})

		When("the security group does not exist", func() {
			It("creates the security group", func() {
				session := helpers.CF("create-security-group", securityGroupName, tempPath)
				Eventually(session).Should(Say("Creating security group %s as %s...", securityGroupName, userName))
				Eventually(session).Should(Say("OK"))
				Eventually(session).Should(Exit(0))

				session = helpers.CF("security-group", securityGroupName)
				Eventually(session).Should(Say(`Name\s+%s`, securityGroupName))
			})
		})

		When("the security group already exists", func() {
			BeforeEach(func() {
				Eventually(helpers.CF("create-security-group", securityGroupName, tempPath)).Should(Exit(0))
			})

			It("notifies the user that the security group already exists and exits 0", func() {
				session := helpers.CF("create-security-group", securityGroupName, tempPath)
				Eventually(session).Should(Say("Creating security group %s as %s...", securityGroupName, userName))
				Eventually(session.Err).Should(Say(`Security group with name '%s' already exists\.`, securityGroupName))
				Eventually(session).Should(Say("OK"))
				Eventually(session).Should(Exit(0))
			})
		})
	})
})
<commit_msg>V7: fix message formatting in security group test<commit_after>package isolated

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	. "code.cloudfoundry.org/cli/cf/util/testhelpers/matchers"
	"code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("create-security-group command", func() {
	var (
		dir               string
		securityGroupName string
		tempPath          string
		userName          string
	)

	BeforeEach(func() {
		userName = helpers.LoginCF()
		securityGroupName = helpers.NewSecurityGroupName()
		dir = helpers.TempDirAbsolutePath("", "security-group")

		tempPath = filepath.Join(dir, "tmpfile.json")
		err := ioutil.WriteFile(tempPath, []byte(`[{
			"protocol": "tcp",
			"destination": "10.0.11.0/24",
			"ports": "80,443",
			"description": "Allow http and https traffic from ZoneA"
		}]`), 0666)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		err := os.RemoveAll(dir)
		Expect(err).NotTo(HaveOccurred())
	})

	Describe("help", func() {
		When("--help flag is set", func() {
			It("appears in cf help -a", func() {
				session := helpers.CF("help", "-a")
				Eventually(session).Should(Exit(0))
				Expect(session).To(HaveCommandInCategoryWithDescription("create-security-group", "SECURITY GROUP", "Create a security group"))
			})

			It("Displays command usage to output", func() {
				session := helpers.CF("create-security-group", "--help")
				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Say("create-security-group - Create a security group"))
				Eventually(session).Should(Say("USAGE:"))
				Eventually(session).Should(Say(`cf create-security-group SECURITY_GROUP PATH_TO_JSON_RULES_FILE`))
				Eventually(session).Should(Say(`The provided path can be an absolute or relative path to a file. The file should have`))
				Eventually(session).Should(Say(`a single array with JSON objects inside describing the rules. The JSON Base Object is`))
				Eventually(session).Should(Say(`omitted and only the square brackets and associated child object are required in the file.`))

				Eventually(session).Should(Say(`Valid json file example:`))
				Eventually(session).Should(Say(`\s+\[`))
				Eventually(session).Should(Say(`\s+{`))
				Eventually(session).Should(Say(`\s+"protocol": "tcp",`))
				Eventually(session).Should(Say(`\s+"destination": "10.0.11.0/24",`))
				Eventually(session).Should(Say(`\s+"ports": "80,443",`))
				Eventually(session).Should(Say(`\s+"description": "Allow http and https traffic from ZoneA"`))
				Eventually(session).Should(Say(`\s+}`))
				Eventually(session).Should(Say(`\s+\]`))
				Eventually(session).Should(Say("SEE ALSO:"))
				Eventually(session).Should(Say("bind-running-security-group, bind-security-group, bind-staging-security-group, security-groups"))
				Eventually(session).Should(Exit(0))
			})
		})
	})

	When("not logged in", func() {
		It("fails with the appropriate errors", func() {
			helpers.CheckEnvironmentTargetedCorrectly(false, false, "", "create-security-group", "some-group", tempPath)
		})
	})

	When("the environment is set up correctly", func() {
		When("the security group name is not provided", func() {
			It("tells the user that the security group name is required, prints help text, and exits 1", func() {
				session := helpers.CF("create-security-group")

				Eventually(session.Err).Should(Say("Incorrect Usage: the required arguments `SECURITY_GROUP` and `PATH_TO_JSON_RULES_FILE` were not provided"))
				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("the security group path is not valid", func() {
			It("tells the user that the security group path does not exist, prints help text, and exits 1", func() {
				session := helpers.CF("create-security-group", securityGroupName, "/invalid/path")

				Eventually(session.Err).Should(Say("Incorrect Usage: The specified path '/invalid/path' does not exist."))
				Eventually(session).Should(Say("NAME:"))
				Eventually(session).Should(Exit(1))
			})
		})

		When("the security group path has invalid JSON", func() {
			BeforeEach(func() {
				err := ioutil.WriteFile(tempPath, []byte("Invalid JSON!"), 0666)
				Expect(err).ToNot(HaveOccurred())
			})

			It("tells the user that the security group path is required, prints help text, and exits 1", func() {
				session := helpers.CF("create-security-group", securityGroupName, tempPath)
				Eventually(session).Should(Say("Creating security group %s as %s...", securityGroupName, userName))
				Eventually(session.Err).Should(Say("Incorrect json format: %s", strings.ReplaceAll(tempPath, "\\", "\\\\")))

				Eventually(session.Err).Should(Say("Valid json file example:"))
				Eventually(session.Err).Should(Say(`\[`))
				Eventually(session.Err).Should(Say(`\s+{`))
				Eventually(session.Err).Should(Say(`\s+"protocol": "tcp",`))
				Eventually(session.Err).Should(Say(`\s+"destination": "10.244.1.18",`))
				Eventually(session.Err).Should(Say(`\s+"ports": "3306"`))
				Eventually(session.Err).Should(Say(`\s+}`))
				Eventually(session.Err).Should(Say(`\]`))
				Eventually(session).Should(Exit(1))
			})
		})

		When("the security group does not exist", func() {
			It("creates the security group", func() {
				session := helpers.CF("create-security-group", securityGroupName, tempPath)
				Eventually(session).Should(Say("Creating security group %s as %s...", securityGroupName, userName))
				Eventually(session).Should(Say("OK"))
				Eventually(session).Should(Exit(0))

				session = helpers.CF("security-group", securityGroupName)
				Eventually(session).Should(Say(`name:\s+%s`, securityGroupName))
			})
		})

		When("the security group already exists", func() {
			BeforeEach(func() {
				Eventually(helpers.CF("create-security-group", securityGroupName, tempPath)).Should(Exit(0))
			})

			It("notifies the user that the security group already exists and exits 0", func() {
				session := helpers.CF("create-security-group", securityGroupName, tempPath)
				Eventually(session).Should(Say("Creating security group %s as %s...", securityGroupName, userName))
				Eventually(session.Err).Should(Say(`Security group with name '%s' already exists\.`, securityGroupName))
				Eventually(session).Should(Say("OK"))
				Eventually(session).Should(Exit(0))
			})
		})
	})
})
<|endoftext|>


<commit_before>package elastic

import (
	"flag"
	"fmt"
	"testing"

	. "gopkg.in/check.v1"
)

var host = flag.String("host", "127.0.0.1", "Elasticsearch host")
var port = flag.Int("port", 9200, "Elasticsearch port")

func Test(t *testing.T) {
	TestingT(t)
}

type elasticTestSuite struct {
	c *Client
}

var _ = Suite(&elasticTestSuite{})

func (s *elasticTestSuite) SetUpSuite(c *C) {
	cfg := new(ClientConfig)
	cfg.Addr = fmt.Sprintf("%s:%d", *host, *port)
	cfg.User = ""
	cfg.Password = ""
	s.c = NewClient(cfg)
}

func (s *elasticTestSuite) TearDownSuite(c *C) {

}

func makeTestData(arg1 string, arg2 string) map[string]interface{} {
	m := make(map[string]interface{})
	m["name"] = arg1
	m["content"] = arg2

	return m
}

func (s *elasticTestSuite) TestSimple(c *C) {
	index := "dummy"
	docType := "blog"

	//key1 := "name"
	//key2 := "content"

	err := s.c.Update(index, docType, "1", makeTestData("abc", "hello world"))
	c.Assert(err, IsNil)

	exists, err := s.c.Exists(index, docType, "1")
	c.Assert(err, IsNil)
	c.Assert(exists, Equals, true)

	r, err := s.c.Get(index, docType, "1")
	c.Assert(err, IsNil)
	c.Assert(r.Code, Equals, 200)
	c.Assert(r.ID, Equals, "1")

	err = s.c.Delete(index, docType, "1")
	c.Assert(err, IsNil)

	exists, err = s.c.Exists(index, docType, "1")
	c.Assert(err, IsNil)
	c.Assert(exists, Equals, false)

	items := make([]*BulkRequest, 10)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Action = ActionIndex
		req.ID = id
		req.Data = makeTestData(fmt.Sprintf("abc %d", i), fmt.Sprintf("hello world %d", i))
		items[i] = req
	}

	resp, err := s.c.IndexTypeBulk(index, docType, items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Action = ActionDelete
		req.ID = id
		items[i] = req
	}

	resp, err = s.c.IndexTypeBulk(index, docType, items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)
}

// this requires a parent setting in _mapping
func (s *elasticTestSuite) TestParent(c *C) {
	index := "dummy"
	docType := "comment"
	ParentType := "parent"

	mapping := map[string]interface{}{
			docType: map[string]interface{}{
				"_parent": map[string]string{"type": ParentType},
		},
	}
	err := s.c.CreateMapping(index, docType,  mapping)
	c.Assert(err, IsNil)

	items := make([]*BulkRequest, 10)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Action = ActionIndex
		req.ID = id
		req.Data = makeTestData(fmt.Sprintf("abc %d", i), fmt.Sprintf("hello world %d", i))
		req.Parent = "1"
		items[i] = req
	}

	resp, err := s.c.IndexTypeBulk(index, docType, items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Index = index
		req.Type = docType
		req.Action = ActionDelete
		req.ID = id
		req.Parent = "1"
		items[i] = req
	}
	resp, err = s.c.Bulk(items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)
}
<commit_msg>Update client_test.go<commit_after>package elastic

import (
	"flag"
	"fmt"
	"testing"

	. "gopkg.in/check.v1"
)

var host = flag.String("host", "127.0.0.1", "Elasticsearch host")
var port = flag.Int("port", 9200, "Elasticsearch port")

func Test(t *testing.T) {
	TestingT(t)
}

type elasticTestSuite struct {
	c *Client
}

var _ = Suite(&elasticTestSuite{})

func (s *elasticTestSuite) SetUpSuite(c *C) {
	cfg := new(ClientConfig)
	cfg.Addr = fmt.Sprintf("%s:%d", *host, *port)
	cfg.User = ""
	cfg.Password = ""
	s.c = NewClient(cfg)
}

func (s *elasticTestSuite) TearDownSuite(c *C) {

}

func makeTestData(arg1 string, arg2 string) map[string]interface{} {
	m := make(map[string]interface{})
	m["name"] = arg1
	m["content"] = arg2

	return m
}

func (s *elasticTestSuite) TestSimple(c *C) {
	index := "dummy"
	docType := "blog"

	//key1 := "name"
	//key2 := "content"

	err := s.c.Update(index, docType, "1", makeTestData("abc", "hello world"))
	c.Assert(err, IsNil)

	exists, err := s.c.Exists(index, docType, "1")
	c.Assert(err, IsNil)
	c.Assert(exists, Equals, true)

	r, err := s.c.Get(index, docType, "1")
	c.Assert(err, IsNil)
	c.Assert(r.Code, Equals, 200)
	c.Assert(r.ID, Equals, "1")

	err = s.c.Delete(index, docType, "1")
	c.Assert(err, IsNil)

	exists, err = s.c.Exists(index, docType, "1")
	c.Assert(err, IsNil)
	c.Assert(exists, Equals, false)

	items := make([]*BulkRequest, 10)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Action = ActionIndex
		req.ID = id
		req.Data = makeTestData(fmt.Sprintf("abc %d", i), fmt.Sprintf("hello world %d", i))
		items[i] = req
	}

	resp, err := s.c.IndexTypeBulk(index, docType, items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Action = ActionDelete
		req.ID = id
		items[i] = req
	}

	resp, err = s.c.IndexTypeBulk(index, docType, items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)
}

// this requires a parent setting in _mapping
func (s *elasticTestSuite) TestParent(c *C) {
	index := "dummy"
	docType := "comment"
	ParentType := "parent"

	mapping := map[string]interface{}{
			docType: map[string]interface{}{
				"_parent": map[string]string{"type": ParentType},
		},
	}
	err := s.c.CreateMapping(index, docType,  mapping)
	c.Assert(err, IsNil)

	items := make([]*BulkRequest, 10)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Action = ActionIndex
		req.ID = id
		req.Data = makeTestData(fmt.Sprintf("abc %d", i), fmt.Sprintf("hello world %d", i))
		req.Parent = "1"
		items[i] = req
	}

	resp, err := s.c.IndexTypeBulk(index, docType, items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)

	for i := 0; i < 10; i++ {
		id := fmt.Sprintf("%d", i)
		req := new(BulkRequest)
		req.Index = index
		req.Type = docType
		req.Action = ActionDelete
		req.ID = id
		req.Parent = "1"
		items[i] = req
	}
	resp, err = s.c.Bulk(items)
	c.Assert(err, IsNil)
	c.Assert(resp.Code, Equals, 200)
	c.Assert(resp.Errors, Equals, false)
}
<|endoftext|>


<commit_before>// The server program issues Google search requests and demonstrates the use of
// the go.net Context API. It serves on port 8080.
//
// The /search endpoint accepts these query params:
//   q=the Google search query
//   timeout=a timeout for the request, in time.Duration format
//
// For example, http://localhost:8080/search?q=golang&timeout=1s serves the
// first few Google search results for "golang" or a "deadline exceeded" error
// if the timeout expires.
package main

import (
	"html/template"
	"log"
	"net/http"
	"time"

	"code.google.com/p/go.net/context"
	"github.com/niilo/golib/context/google"
	"github.com/niilo/golib/context/userip"
)

func main() {
	http.HandleFunc("/search", handleSearch)
	log.Fatal(http.ListenAndServe(":8080", nil))
}

// handleSearch handles URLs like /search?q=golang&timeout=1s by forwarding the
// query to google.Search. If the query param includes timeout, the search is
// canceled after that duration elapses.
func handleSearch(w http.ResponseWriter, req *http.Request) {
	// ctx is the Context for this handler. Calling cancel closes the
	// ctx.Done channel, which is the cancellation signal for requests
	// started by this handler.
	var (
		ctx    context.Context
		cancel context.CancelFunc
	)
	timeout, err := time.ParseDuration(req.FormValue("timeout"))
	if err == nil {
		// The request has a timeout, so create a context that is
		// canceled automatically when the timeout expires.
		ctx, cancel = context.WithTimeout(context.Background(), timeout)
	} else {
		ctx, cancel = context.WithCancel(context.Background())
	}
	defer cancel() // Cancel ctx as soon as handleSearch returns.

	// Check the search query.
	query := req.FormValue("q")
	if query == "" {
		http.Error(w, "no query", http.StatusBadRequest)
		return
	}

	// Store the user IP in ctx for use by code in other packages.
	userIP, err := userip.FromRequest(req)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	ctx = userip.NewContext(ctx, userIP)

	// Run the Google search and print the results.
	start := time.Now()
	results, err := google.Search(ctx, query)
	elapsed := time.Since(start)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if err := resultsTemplate.Execute(w, struct {
		Results          google.Results
		Timeout, Elapsed time.Duration
	}{
		Results: results,
		Timeout: timeout,
		Elapsed: elapsed,
	}); err != nil {
		log.Print(err)
		return
	}
}

var resultsTemplate = template.Must(template.New("results").Parse(`

  {{range .Results}}

{{.Title}} - {{.URL}}

  {{end}}

{{len .Results}} results in {{.Elapsed}}; timeout {{.Timeout}}

`))
<commit_msg>add httprouter and alice middleware<commit_after>// The server program issues Google search requests and demonstrates the use of
// the go.net Context API. It serves on port 8080.
//
// The /search endpoint accepts these query params:
//   q=the Google search query
//   timeout=a timeout for the request, in time.Duration format
//
// For example, http://localhost:8080/search?q=golang&timeout=1s serves the
// first few Google search results for "golang" or a "deadline exceeded" error
// if the timeout expires.
package main

import (
	"fmt"
	"html/template"
	"log"
	"net/http"
	"time"

	"code.google.com/p/go.net/context"
	"github.com/PuerkitoBio/throttled"
	"github.com/julienschmidt/httprouter"
	"github.com/justinas/alice"
	"github.com/niilo/golib/context/google"
	"github.com/niilo/golib/context/userip"
)

func Index(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	fmt.Fprint(w, "Welcome!\n")
}

func Hello(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	fmt.Fprintf(w, "hello, %s!\n", ps.ByName("name"))
}

func timeoutHandler(h http.Handler) http.Handler {
	return http.TimeoutHandler(h, 1*time.Second, "timed out")
}

//func myApp(w http.ResponseWriter, r *http.Request) {
//	w.Write([]byte("Hello world!"))
//}

func main() {
	router := httprouter.New()
	router.GET("/", Index)
	router.GET("/hello/:name", Hello)
	router.GET("/search", handleSearch)

	//log.Fatal(http.ListenAndServe(":8080", router))

	th := throttled.Interval(throttled.PerSec(10), 1, &throttled.VaryBy{Path: true}, 50)
	//myHandler := http.HandlerFunc(myApp)

	chain := alice.New(th.Throttle, timeoutHandler).Then(router)
	http.ListenAndServe(":8080", chain)

}

// handleSearch handles URLs like /search?q=golang&timeout=1s by forwarding the
// query to google.Search. If the query param includes timeout, the search is
// canceled after that duration elapses.
func handleSearch(w http.ResponseWriter, req *http.Request, _ httprouter.Params) {
	// ctx is the Context for this handler. Calling cancel closes the
	// ctx.Done channel, which is the cancellation signal for requests
	// started by this handler.
	var (
		ctx    context.Context
		cancel context.CancelFunc
	)
	timeout, err := time.ParseDuration(req.FormValue("timeout"))
	if err == nil {
		// The request has a timeout, so create a context that is
		// canceled automatically when the timeout expires.
		ctx, cancel = context.WithTimeout(context.Background(), timeout)
	} else {
		ctx, cancel = context.WithCancel(context.Background())
	}
	defer cancel() // Cancel ctx as soon as handleSearch returns.

	// Check the search query.
	query := req.FormValue("q")
	if query == "" {
		http.Error(w, "no query", http.StatusBadRequest)
		return
	}

	// Store the user IP in ctx for use by code in other packages.
	userIP, err := userip.FromRequest(req)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	ctx = userip.NewContext(ctx, userIP)

	// Run the Google search and print the results.
	start := time.Now()
	results, err := google.Search(ctx, query)
	elapsed := time.Since(start)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if err := resultsTemplate.Execute(w, struct {
		Results          google.Results
		Timeout, Elapsed time.Duration
	}{
		Results: results,
		Timeout: timeout,
		Elapsed: elapsed,
	}); err != nil {
		log.Print(err)
		return
	}
}

var resultsTemplate = template.Must(template.New("results").Parse(`

  {{range .Results}}

{{.Title}} - {{.URL}}

  {{end}}

{{len .Results}} results in {{.Elapsed}}; timeout {{.Timeout}}

`))
<|endoftext|>


<commit_before>package lru

import (
	"bytes"
	"container/list"
	"encoding/base64"
	"fmt"
	"hash/fnv"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"time"
)

type object struct {
	key        string
	size       uint64
	accessTime time.Time
}

type Cache struct {
	sync.Mutex

	list  *list.List
	table map[string]*list.Element
	size  uint64

	maxSize uint64
	path    string

	debug bool
}

type Options struct {
	ClearCacheOnBoot bool
	Debug            bool
}

func hashCacheKey(data string) string {
	hash := fnv.New64a()
	hash.Write([]byte(data))

	return base64.URLEncoding.EncodeToString(hash.Sum(nil))
}

func New(maxSize uint64, path string, options Options) *Cache {
	cache := Cache{
		list:    list.New(),
		table:   make(map[string]*list.Element),
		maxSize: maxSize,
		path:    path,
		debug:   options.Debug,
	}

	cache.Debug(fmt.Sprintf("new cache of size %d", maxSize))

	if options.ClearCacheOnBoot {
		cache.Debug("clearing cache on boot")
		os.RemoveAll(cache.path)
		os.MkdirAll(cache.path, 0755)
	}

	return &cache
}

func (cache *Cache) Debug(msg string) {
	if cache.debug {
		fmt.Println("[lru]", msg)
	}
}

func (cache *Cache) FilePath(key string) string {
	return cache.path + "/" + hashCacheKey(key)
}

func (cache *Cache) Get(key string) ([]byte, bool) {
	cache.Lock()
	defer cache.Unlock()

	element := cache.table[key]

	if element == nil {
		return nil, false
	}

	cache.moveToFront(element)

	value, err := ioutil.ReadFile(cache.FilePath(element.Value.(*object).key))

	if err != nil {
		return nil, false
	}

	return value, true
}

func (cache *Cache) GetBuffer(key string) (data *bytes.Buffer, ok bool) {
	cache.Lock()
	defer cache.Unlock()

	element := cache.table[key]

	if element == nil {
		return nil, false
	}

	cache.moveToFront(element)

	file, err := os.Open(cache.FilePath(element.Value.(*object).key))

	if err != nil {
		return nil, false
	}

	data = &bytes.Buffer{}

	io.Copy(data, file)

	if err != nil {
		return nil, false
	}

	return data, true
}

func (cache *Cache) Set(key string, value []byte) {
	cache.Lock()
	defer cache.Unlock()

	if element := cache.table[key]; element != nil {
		cache.moveToFront(element)
	} else {
		cache.addNew(key, value)
	}
}

func (cache *Cache) Delete(key string) bool {
	cache.Lock()
	defer cache.Unlock()

	element := cache.table[key]

	if element == nil {
		return false
	}

	err := os.Remove(cache.FilePath(key))

	if err != nil {
		return false
	}

	cache.list.Remove(element)
	delete(cache.table, key)

	cache.size -= element.Value.(*object).size

	return true
}

func (cache *Cache) Clear() {
	cache.Lock()
	defer cache.Unlock()

	cache.clearFiles()
	cache.list.Init()
	cache.table = make(map[string]*list.Element)
	cache.size = 0
}

func (cache *Cache) Size() uint64 {
	cache.Lock()
	defer cache.Unlock()

	return cache.size
}

func (cache *Cache) MaxSize() uint64 {
	cache.Lock()
	defer cache.Unlock()

	return cache.maxSize
}

func (cache *Cache) Oldest() (oldest time.Time) {
	cache.Lock()
	defer cache.Unlock()

	if lastElem := cache.list.Back(); lastElem != nil {
		oldest = lastElem.Value.(*object).accessTime
	}

	return
}

func (cache *Cache) keys() []string {
	keys := make([]string, 0, cache.list.Len())

	for element := cache.list.Front(); element != nil; element = element.Next() {
		keys = append(keys, element.Value.(*object).key)
	}

	return keys
}

func (cache *Cache) moveToFront(element *list.Element) {
	cache.list.MoveToFront(element)
	element.Value.(*object).accessTime = time.Now()
}

func (cache *Cache) addNew(key string, value []byte) {
	size := uint64(len(value))

	cache.Debug(fmt.Sprintf("new object of size %d", size))

	if size > cache.maxSize {
		cache.Debug("file is too large")
		return
	}

	newObject := &object{key, size, time.Now()}

	cache.trim(cache.size + newObject.size)

	if _, err := os.Stat(cache.FilePath(key)); os.IsNotExist(err) {
		err := ioutil.WriteFile(cache.FilePath(key), value, 0644)

		if err != nil {
			cache.Debug(err.Error())
			return
		}

		element := cache.list.PushFront(newObject)
		cache.table[key] = element
		cache.size += (*newObject).size
		cache.Debug(fmt.Sprintf("added %d, new size is %d", (*newObject).size, cache.size))
	} else {
		cache.Debug("file already exist")
	}
}

func (cache *Cache) trim(futureSize uint64) {
	for futureSize > cache.maxSize {
		element := cache.list.Back()

		if element == nil {
			cache.Debug("file is too large")
			return
		}

		value := cache.list.Remove(element).(*object)

		cache.Debug(fmt.Sprintf("deleting %s", cache.FilePath(value.key)))

		if err := os.RemoveAll(cache.FilePath(value.key)); err != nil {
			cache.Debug(fmt.Sprintf("couldn't delete %s", cache.FilePath(value.key)))
		}

		delete(cache.table, value.key)

		cache.size -= value.size
		futureSize -= value.size
	}
}

func (cache *Cache) clearFiles() {
	for _, key := range cache.keys() {
		os.RemoveAll(cache.FilePath(key))
	}
}
<commit_msg>Use a buffer pool to reduce allocation<commit_after>package lru

import (
	"bytes"
	"container/list"
	"encoding/base64"
	"fmt"
	"hash/fnv"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"time"
)

var (
	bufferPool = sync.Pool{
		New: func() interface{} { return &bytes.Buffer{} },
	}
)

type object struct {
	key        string
	size       uint64
	accessTime time.Time
}

type Cache struct {
	sync.Mutex

	list  *list.List
	table map[string]*list.Element
	size  uint64

	maxSize uint64
	path    string

	debug bool
}

type Options struct {
	ClearCacheOnBoot bool
	Debug            bool
}

func hashCacheKey(data string) string {
	hash := fnv.New64a()
	hash.Write([]byte(data))

	return base64.URLEncoding.EncodeToString(hash.Sum(nil))
}

func New(maxSize uint64, path string, options Options) *Cache {
	cache := Cache{
		list:    list.New(),
		table:   make(map[string]*list.Element),
		maxSize: maxSize,
		path:    path,
		debug:   options.Debug,
	}

	cache.Debug(fmt.Sprintf("new cache of size %d", maxSize))

	if options.ClearCacheOnBoot {
		cache.Debug("clearing cache on boot")
		os.RemoveAll(cache.path)
		os.MkdirAll(cache.path, 0755)
	}

	return &cache
}

func (cache *Cache) Debug(msg string) {
	if cache.debug {
		fmt.Println("[lru]", msg)
	}
}

func (cache *Cache) FilePath(key string) string {
	return cache.path + "/" + hashCacheKey(key)
}

func (cache *Cache) Get(key string) ([]byte, bool) {
	cache.Lock()
	defer cache.Unlock()

	element := cache.table[key]

	if element == nil {
		return nil, false
	}

	cache.moveToFront(element)

	value, err := ioutil.ReadFile(cache.FilePath(element.Value.(*object).key))

	if err != nil {
		return nil, false
	}

	return value, true
}

func (cache *Cache) GetBuffer(key string) (data *bytes.Buffer, ok bool) {
	cache.Lock()
	defer cache.Unlock()

	element := cache.table[key]

	if element == nil {
		return nil, false
	}

	cache.moveToFront(element)

	file, err := os.Open(cache.FilePath(element.Value.(*object).key))

	if err != nil {
		return nil, false
	}

	data = bufferPool.Get().(*bytes.Buffer)

	data.Reset()
	io.Copy(data, file)

	if err != nil {
		return nil, false
	}

	return data, true
}

func (cache *Cache) Set(key string, value []byte) {
	cache.Lock()
	defer cache.Unlock()

	if element := cache.table[key]; element != nil {
		cache.moveToFront(element)
	} else {
		cache.addNew(key, value)
	}
}

func (cache *Cache) Delete(key string) bool {
	cache.Lock()
	defer cache.Unlock()

	element := cache.table[key]

	if element == nil {
		return false
	}

	err := os.Remove(cache.FilePath(key))

	if err != nil {
		return false
	}

	cache.list.Remove(element)
	delete(cache.table, key)

	cache.size -= element.Value.(*object).size

	return true
}

func (cache *Cache) Clear() {
	cache.Lock()
	defer cache.Unlock()

	cache.clearFiles()
	cache.list.Init()
	cache.table = make(map[string]*list.Element)
	cache.size = 0
}

func (cache *Cache) Size() uint64 {
	cache.Lock()
	defer cache.Unlock()

	return cache.size
}

func (cache *Cache) MaxSize() uint64 {
	cache.Lock()
	defer cache.Unlock()

	return cache.maxSize
}

func (cache *Cache) Oldest() (oldest time.Time) {
	cache.Lock()
	defer cache.Unlock()

	if lastElem := cache.list.Back(); lastElem != nil {
		oldest = lastElem.Value.(*object).accessTime
	}

	return
}

func (cache *Cache) keys() []string {
	keys := make([]string, 0, cache.list.Len())

	for element := cache.list.Front(); element != nil; element = element.Next() {
		keys = append(keys, element.Value.(*object).key)
	}

	return keys
}

func (cache *Cache) moveToFront(element *list.Element) {
	cache.list.MoveToFront(element)
	element.Value.(*object).accessTime = time.Now()
}

func (cache *Cache) addNew(key string, value []byte) {
	size := uint64(len(value))

	cache.Debug(fmt.Sprintf("new object of size %d", size))

	if size > cache.maxSize {
		cache.Debug("file is too large")
		return
	}

	newObject := &object{key, size, time.Now()}

	cache.trim(cache.size + newObject.size)

	if _, err := os.Stat(cache.FilePath(key)); os.IsNotExist(err) {
		err := ioutil.WriteFile(cache.FilePath(key), value, 0644)

		if err != nil {
			cache.Debug(err.Error())
			return
		}

		element := cache.list.PushFront(newObject)
		cache.table[key] = element
		cache.size += (*newObject).size
		cache.Debug(fmt.Sprintf("added %d, new size is %d", (*newObject).size, cache.size))
	} else {
		cache.Debug("file already exist")
	}
}

func (cache *Cache) trim(futureSize uint64) {
	for futureSize > cache.maxSize {
		element := cache.list.Back()

		if element == nil {
			cache.Debug("file is too large")
			return
		}

		value := cache.list.Remove(element).(*object)

		cache.Debug(fmt.Sprintf("deleting %s", cache.FilePath(value.key)))

		if err := os.RemoveAll(cache.FilePath(value.key)); err != nil {
			cache.Debug(fmt.Sprintf("couldn't delete %s", cache.FilePath(value.key)))
		}

		delete(cache.table, value.key)

		cache.size -= value.size
		futureSize -= value.size
	}
}

func (cache *Cache) clearFiles() {
	for _, key := range cache.keys() {
		os.RemoveAll(cache.FilePath(key))
	}
}
<|endoftext|>


<commit_before>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sqlparser

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/wandoulabs/cm/hack"
	"github.com/wandoulabs/cm/sqltypes"
)

const EOFCHAR = 0x100

// Tokenizer is the struct used to generate SQL
// tokens for the parser.
type Tokenizer struct {
	InStream      *strings.Reader
	AllowComments bool
	ForceEOF      bool
	lastChar      uint16
	Position      int
	errorToken    []byte
	LastError     string
	posVarIndex   int
	ParseTree     Statement
}

// NewStringTokenizer creates a new Tokenizer for the
// sql string.
func NewStringTokenizer(sql string) *Tokenizer {
	return &Tokenizer{InStream: strings.NewReader(sql)}
}

var keywords = map[string]int{
	"all":           ALL,
	"alter":         ALTER,
	"analyze":       ANALYZE,
	"and":           AND,
	"as":            AS,
	"asc":           ASC,
	"between":       BETWEEN,
	"by":            BY,
	"case":          CASE,
	"create":        CREATE,
	"cross":         CROSS,
	"default":       DEFAULT,
	"delete":        DELETE,
	"desc":          DESC,
	"describe":      DESCRIBE,
	"distinct":      DISTINCT,
	"drop":          DROP,
	"duplicate":     DUPLICATE,
	"else":          ELSE,
	"end":           END,
	"except":        EXCEPT,
	"exists":        EXISTS,
	"explain":       EXPLAIN,
	"for":           FOR,
	"force":         FORCE,
	"from":          FROM,
	"group":         GROUP,
	"having":        HAVING,
	"if":            IF,
	"ignore":        IGNORE,
	"in":            IN,
	"index":         INDEX,
	"inner":         INNER,
	"insert":        INSERT,
	"intersect":     INTERSECT,
	"into":          INTO,
	"is":            IS,
	"join":          JOIN,
	"key":           KEY,
	"left":          LEFT,
	"like":          LIKE,
	"limit":         LIMIT,
	"lock":          LOCK,
	"minus":         MINUS,
	"natural":       NATURAL,
	"not":           NOT,
	"null":          NULL,
	"on":            ON,
	"or":            OR,
	"order":         ORDER,
	"outer":         OUTER,
	"rename":        RENAME,
	"right":         RIGHT,
	"select":        SELECT,
	"set":           SET,
	"show":          SHOW,
	"straight_join": STRAIGHT_JOIN,
	"table":         TABLE,
	"then":          THEN,
	"to":            TO,
	"union":         UNION,
	"unique":        UNIQUE,
	"update":        UPDATE,
	"use":           USE,
	"using":         USING,
	"values":        VALUES,
	"view":          VIEW,
	"when":          WHEN,
	"where":         WHERE,
}

// Lex returns the next token form the Tokenizer.
// This function is used by go yacc.
func (tkn *Tokenizer) Lex(lval *yySymType) int {
	typ, val := tkn.Scan()
	for typ == COMMENT {
		if tkn.AllowComments {
			break
		}
		typ, val = tkn.Scan()
	}
	switch typ {
	case ID, STRING, NUMBER, VALUE_ARG, LIST_ARG, COMMENT:
		lval.bytes = val
	}
	tkn.errorToken = val
	return typ
}

// Error is called by go yacc if there's a parsing error.
func (tkn *Tokenizer) Error(err string) {
	buf := bytes.NewBuffer(make([]byte, 0, 32))
	if tkn.errorToken != nil {
		fmt.Fprintf(buf, "%s at position %v near %s", err, tkn.Position, tkn.errorToken)
	} else {
		fmt.Fprintf(buf, "%s at position %v", err, tkn.Position)
	}
	tkn.LastError = buf.String()
}

// Scan scans the tokenizer for the next token and returns
// the token type and an optional value.
func (tkn *Tokenizer) Scan() (int, []byte) {
	if tkn.ForceEOF {
		return 0, nil
	}

	if tkn.lastChar == 0 {
		tkn.next()
	}
	tkn.skipBlank()
	switch ch := tkn.lastChar; {
	case isLetter(ch):
		return tkn.scanIdentifier()
	case isDigit(ch):
		return tkn.scanNumber(false)
	case ch == ':':
		return tkn.scanBindVar()
	default:
		tkn.next()
		switch ch {
		case EOFCHAR:
			return 0, nil
		case '=', ',', ';', '(', ')', '+', '*', '%', '&', '|', '^', '~':
			return int(ch), nil
		case '?':
			tkn.posVarIndex++
			buf := new(bytes.Buffer)
			fmt.Fprintf(buf, ":v%d", tkn.posVarIndex)
			return VALUE_ARG, buf.Bytes()
		case '.':
			if isDigit(tkn.lastChar) {
				return tkn.scanNumber(true)
			} else {
				return int(ch), nil
			}
		case '/':
			switch tkn.lastChar {
			case '/':
				tkn.next()
				return tkn.scanCommentType1("//")
			case '*':
				tkn.next()
				return tkn.scanCommentType2()
			default:
				return int(ch), nil
			}
		case '-':
			if tkn.lastChar == '-' {
				tkn.next()
				return tkn.scanCommentType1("--")
			} else {
				return int(ch), nil
			}
		case '<':
			switch tkn.lastChar {
			case '>':
				tkn.next()
				return NE, nil
			case '=':
				tkn.next()
				switch tkn.lastChar {
				case '>':
					tkn.next()
					return NULL_SAFE_EQUAL, nil
				default:
					return LE, nil
				}
			default:
				return int(ch), nil
			}
		case '>':
			if tkn.lastChar == '=' {
				tkn.next()
				return GE, nil
			} else {
				return int(ch), nil
			}
		case '!':
			if tkn.lastChar == '=' {
				tkn.next()
				return NE, nil
			} else {
				return LEX_ERROR, []byte("!")
			}
		case '\'', '"':
			return tkn.scanString(ch, STRING)
		case '`':
			return tkn.scanLiteralIdentifier()
		default:
			return LEX_ERROR, []byte{byte(ch)}
		}
	}
}

func (tkn *Tokenizer) skipBlank() {
	ch := tkn.lastChar
	for ch == ' ' || ch == '\n' || ch == '\r' || ch == '\t' {
		tkn.next()
		ch = tkn.lastChar
	}
}

func (tkn *Tokenizer) scanIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	buffer.WriteByte(byte(tkn.lastChar))
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	lowered := bytes.ToLower(buffer.Bytes())
	if keywordId, found := keywords[hack.String(lowered)]; found {
		return keywordId, lowered
	}

	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanLiteralIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	if tkn.lastChar != '`' {
		return LEX_ERROR, buffer.Bytes()
	}
	tkn.next()
	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanBindVar() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	token := VALUE_ARG
	tkn.next()
	if tkn.lastChar == ':' {
		token = LIST_ARG
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for isLetter(tkn.lastChar) || isDigit(tkn.lastChar) || tkn.lastChar == '.' {
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	return token, buffer.Bytes()
}

func (tkn *Tokenizer) scanMantissa(base int, buffer *bytes.Buffer) {
	for digitVal(tkn.lastChar) < base {
		tkn.ConsumeNext(buffer)
	}
}

func (tkn *Tokenizer) scanNumber(seenDecimalPoint bool) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	if seenDecimalPoint {
		buffer.WriteByte('.')
		tkn.scanMantissa(10, buffer)
		goto exponent
	}

	if tkn.lastChar == '0' {
		// int or float
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == 'x' || tkn.lastChar == 'X' {
			// hexadecimal int
			tkn.ConsumeNext(buffer)
			tkn.scanMantissa(16, buffer)
		} else {
			// octal int or float
			seenDecimalDigit := false
			tkn.scanMantissa(8, buffer)
			if tkn.lastChar == '8' || tkn.lastChar == '9' {
				// illegal octal int or float
				seenDecimalDigit = true
				tkn.scanMantissa(10, buffer)
			}
			if tkn.lastChar == '.' || tkn.lastChar == 'e' || tkn.lastChar == 'E' {
				goto fraction
			}
			// octal int
			if seenDecimalDigit {
				return LEX_ERROR, buffer.Bytes()
			}
		}
		goto exit
	}

	// decimal int or float
	tkn.scanMantissa(10, buffer)

fraction:
	if tkn.lastChar == '.' {
		tkn.ConsumeNext(buffer)
		tkn.scanMantissa(10, buffer)
	}

exponent:
	if tkn.lastChar == 'e' || tkn.lastChar == 'E' {
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == '+' || tkn.lastChar == '-' {
			tkn.ConsumeNext(buffer)
		}
		tkn.scanMantissa(10, buffer)
	}

exit:
	return NUMBER, buffer.Bytes()
}

func (tkn *Tokenizer) scanString(delim uint16, typ int) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	for {
		ch := tkn.lastChar
		tkn.next()
		if ch == delim {
			if tkn.lastChar == delim {
				tkn.next()
			} else {
				break
			}
		} else if ch == '\\' {
			if tkn.lastChar == EOFCHAR {
				return LEX_ERROR, buffer.Bytes()
			}
			if decodedChar := sqltypes.SqlDecodeMap[byte(tkn.lastChar)]; decodedChar == sqltypes.DONTESCAPE {
				ch = tkn.lastChar
			} else {
				ch = uint16(decodedChar)
			}
			tkn.next()
		}
		if ch == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		buffer.WriteByte(byte(ch))
	}
	return typ, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType1(prefix string) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString(prefix)
	for tkn.lastChar != EOFCHAR {
		if tkn.lastChar == '\n' {
			tkn.ConsumeNext(buffer)
			break
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType2() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString("/*")
	for {
		if tkn.lastChar == '*' {
			tkn.ConsumeNext(buffer)
			if tkn.lastChar == '/' {
				tkn.ConsumeNext(buffer)
				break
			}
			continue
		}
		if tkn.lastChar == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) ConsumeNext(buffer *bytes.Buffer) {
	if tkn.lastChar == EOFCHAR {
		// This should never happen.
		panic("unexpected EOF")
	}
	buffer.WriteByte(byte(tkn.lastChar))
	tkn.next()
}

func (tkn *Tokenizer) next() {
	if ch, err := tkn.InStream.ReadByte(); err != nil {
		// Only EOF is possible.
		tkn.lastChar = EOFCHAR
	} else {
		tkn.lastChar = uint16(ch)
	}
	tkn.Position++
}

func isLetter(ch uint16) bool {
	return 'a' <= ch && ch <= 'z' || 'A' <= ch && ch <= 'Z' || ch == '_' || ch == '@'
}

func digitVal(ch uint16) int {
	switch {
	case '0' <= ch && ch <= '9':
		return int(ch) - '0'
	case 'a' <= ch && ch <= 'f':
		return int(ch) - 'a' + 10
	case 'A' <= ch && ch <= 'F':
		return int(ch) - 'A' + 10
	}
	return 16 // larger than any legal digit val
}

func isDigit(ch uint16) bool {
	return '0' <= ch && ch <= '9'
}
<commit_msg>hex value support<commit_after>// Copyright 2012, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sqlparser

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/wandoulabs/cm/hack"
	"github.com/wandoulabs/cm/sqltypes"
)

const EOFCHAR = 0x100

// Tokenizer is the struct used to generate SQL
// tokens for the parser.
type Tokenizer struct {
	InStream      *strings.Reader
	AllowComments bool
	ForceEOF      bool
	lastChar      uint16
	Position      int
	errorToken    []byte
	LastError     string
	posVarIndex   int
	ParseTree     Statement
}

// NewStringTokenizer creates a new Tokenizer for the
// sql string.
func NewStringTokenizer(sql string) *Tokenizer {
	return &Tokenizer{InStream: strings.NewReader(sql)}
}

var keywords = map[string]int{
	"all":           ALL,
	"alter":         ALTER,
	"analyze":       ANALYZE,
	"and":           AND,
	"as":            AS,
	"asc":           ASC,
	"between":       BETWEEN,
	"by":            BY,
	"case":          CASE,
	"create":        CREATE,
	"cross":         CROSS,
	"default":       DEFAULT,
	"delete":        DELETE,
	"desc":          DESC,
	"describe":      DESCRIBE,
	"distinct":      DISTINCT,
	"drop":          DROP,
	"duplicate":     DUPLICATE,
	"else":          ELSE,
	"end":           END,
	"except":        EXCEPT,
	"exists":        EXISTS,
	"explain":       EXPLAIN,
	"for":           FOR,
	"force":         FORCE,
	"from":          FROM,
	"group":         GROUP,
	"having":        HAVING,
	"if":            IF,
	"ignore":        IGNORE,
	"in":            IN,
	"index":         INDEX,
	"inner":         INNER,
	"insert":        INSERT,
	"intersect":     INTERSECT,
	"into":          INTO,
	"is":            IS,
	"join":          JOIN,
	"key":           KEY,
	"left":          LEFT,
	"like":          LIKE,
	"limit":         LIMIT,
	"lock":          LOCK,
	"minus":         MINUS,
	"natural":       NATURAL,
	"not":           NOT,
	"null":          NULL,
	"on":            ON,
	"or":            OR,
	"order":         ORDER,
	"outer":         OUTER,
	"rename":        RENAME,
	"right":         RIGHT,
	"select":        SELECT,
	"set":           SET,
	"show":          SHOW,
	"straight_join": STRAIGHT_JOIN,
	"table":         TABLE,
	"then":          THEN,
	"to":            TO,
	"union":         UNION,
	"unique":        UNIQUE,
	"update":        UPDATE,
	"use":           USE,
	"using":         USING,
	"values":        VALUES,
	"view":          VIEW,
	"when":          WHEN,
	"where":         WHERE,
}

// Lex returns the next token form the Tokenizer.
// This function is used by go yacc.
func (tkn *Tokenizer) Lex(lval *yySymType) int {
	typ, val := tkn.Scan()
	for typ == COMMENT {
		if tkn.AllowComments {
			break
		}
		typ, val = tkn.Scan()
	}
	switch typ {
	case ID, STRING, NUMBER, VALUE_ARG, LIST_ARG, COMMENT:
		lval.bytes = val
	}
	tkn.errorToken = val
	return typ
}

// Error is called by go yacc if there's a parsing error.
func (tkn *Tokenizer) Error(err string) {
	buf := bytes.NewBuffer(make([]byte, 0, 32))
	if tkn.errorToken != nil {
		fmt.Fprintf(buf, "%s at position %v near %s", err, tkn.Position, tkn.errorToken)
	} else {
		fmt.Fprintf(buf, "%s at position %v", err, tkn.Position)
	}
	tkn.LastError = buf.String()
}

func (tkn *Tokenizer) scanHexValue() (int, []byte) {
	return 0, nil
}

// Scan scans the tokenizer for the next token and returns
// the token type and an optional value.
func (tkn *Tokenizer) Scan() (int, []byte) {
	if tkn.ForceEOF {
		return 0, nil
	}

	if tkn.lastChar == 0 {
		tkn.next()
	}
	tkn.skipBlank()
	switch ch := tkn.lastChar; {
	case isLetter(ch):
		if ch == 'x' {
			tkn.next()
			c := tkn.lastChar
			if c == '\'' {
				return tkn.scanHexValue()
			}
		}
		return tkn.scanIdentifier()
	case isDigit(ch):
		return tkn.scanNumber(false)
	case ch == ':':
		return tkn.scanBindVar()
	default:
		tkn.next()
		switch ch {
		case EOFCHAR:
			return 0, nil
		case '=', ',', ';', '(', ')', '+', '*', '%', '&', '|', '^', '~':
			return int(ch), nil
		case '?':
			tkn.posVarIndex++
			buf := new(bytes.Buffer)
			fmt.Fprintf(buf, ":v%d", tkn.posVarIndex)
			return VALUE_ARG, buf.Bytes()
		case '.':
			if isDigit(tkn.lastChar) {
				return tkn.scanNumber(true)
			} else {
				return int(ch), nil
			}
		case '/':
			switch tkn.lastChar {
			case '/':
				tkn.next()
				return tkn.scanCommentType1("//")
			case '*':
				tkn.next()
				return tkn.scanCommentType2()
			default:
				return int(ch), nil
			}
		case '-':
			if tkn.lastChar == '-' {
				tkn.next()
				return tkn.scanCommentType1("--")
			} else {
				return int(ch), nil
			}
		case '<':
			switch tkn.lastChar {
			case '>':
				tkn.next()
				return NE, nil
			case '=':
				tkn.next()
				switch tkn.lastChar {
				case '>':
					tkn.next()
					return NULL_SAFE_EQUAL, nil
				default:
					return LE, nil
				}
			default:
				return int(ch), nil
			}
		case '>':
			if tkn.lastChar == '=' {
				tkn.next()
				return GE, nil
			} else {
				return int(ch), nil
			}
		case '!':
			if tkn.lastChar == '=' {
				tkn.next()
				return NE, nil
			} else {
				return LEX_ERROR, []byte("!")
			}
		case '\'', '"':
			return tkn.scanString(ch, STRING)
		case '`':
			return tkn.scanLiteralIdentifier()
		default:
			return LEX_ERROR, []byte{byte(ch)}
		}
	}
}

func (tkn *Tokenizer) skipBlank() {
	ch := tkn.lastChar
	for ch == ' ' || ch == '\n' || ch == '\r' || ch == '\t' {
		tkn.next()
		ch = tkn.lastChar
	}
}

func (tkn *Tokenizer) scanIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	buffer.WriteByte(byte(tkn.lastChar))
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	lowered := bytes.ToLower(buffer.Bytes())
	if keywordId, found := keywords[hack.String(lowered)]; found {
		return keywordId, lowered
	}

	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanLiteralIdentifier() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for tkn.next(); isLetter(tkn.lastChar) || isDigit(tkn.lastChar); tkn.next() {
		buffer.WriteByte(byte(tkn.lastChar))
	}
	if tkn.lastChar != '`' {
		return LEX_ERROR, buffer.Bytes()
	}
	tkn.next()
	return ID, buffer.Bytes()
}

func (tkn *Tokenizer) scanBindVar() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteByte(byte(tkn.lastChar))
	token := VALUE_ARG
	tkn.next()
	if tkn.lastChar == ':' {
		token = LIST_ARG
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	if !isLetter(tkn.lastChar) {
		return LEX_ERROR, buffer.Bytes()
	}
	for isLetter(tkn.lastChar) || isDigit(tkn.lastChar) || tkn.lastChar == '.' {
		buffer.WriteByte(byte(tkn.lastChar))
		tkn.next()
	}
	return token, buffer.Bytes()
}

func (tkn *Tokenizer) scanMantissa(base int, buffer *bytes.Buffer) {
	for digitVal(tkn.lastChar) < base {
		tkn.ConsumeNext(buffer)
	}
}

func (tkn *Tokenizer) scanNumber(seenDecimalPoint bool) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	if seenDecimalPoint {
		buffer.WriteByte('.')
		tkn.scanMantissa(10, buffer)
		goto exponent
	}

	if tkn.lastChar == '0' {
		// int or float
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == 'x' || tkn.lastChar == 'X' {
			// hexadecimal int
			tkn.ConsumeNext(buffer)
			tkn.scanMantissa(16, buffer)
		} else {
			// octal int or float
			seenDecimalDigit := false
			tkn.scanMantissa(8, buffer)
			if tkn.lastChar == '8' || tkn.lastChar == '9' {
				// illegal octal int or float
				seenDecimalDigit = true
				tkn.scanMantissa(10, buffer)
			}
			if tkn.lastChar == '.' || tkn.lastChar == 'e' || tkn.lastChar == 'E' {
				goto fraction
			}
			// octal int
			if seenDecimalDigit {
				return LEX_ERROR, buffer.Bytes()
			}
		}
		goto exit
	}

	// decimal int or float
	tkn.scanMantissa(10, buffer)

fraction:
	if tkn.lastChar == '.' {
		tkn.ConsumeNext(buffer)
		tkn.scanMantissa(10, buffer)
	}

exponent:
	if tkn.lastChar == 'e' || tkn.lastChar == 'E' {
		tkn.ConsumeNext(buffer)
		if tkn.lastChar == '+' || tkn.lastChar == '-' {
			tkn.ConsumeNext(buffer)
		}
		tkn.scanMantissa(10, buffer)
	}

exit:
	return NUMBER, buffer.Bytes()
}

func (tkn *Tokenizer) scanString(delim uint16, typ int) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 64))
	for {
		ch := tkn.lastChar
		tkn.next()
		if ch == delim {
			if tkn.lastChar == delim {
				tkn.next()
			} else {
				break
			}
		} else if ch == '\\' {
			if tkn.lastChar == EOFCHAR {
				return LEX_ERROR, buffer.Bytes()
			}
			if decodedChar := sqltypes.SqlDecodeMap[byte(tkn.lastChar)]; decodedChar == sqltypes.DONTESCAPE {
				ch = tkn.lastChar
			} else {
				ch = uint16(decodedChar)
			}
			tkn.next()
		}
		if ch == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		buffer.WriteByte(byte(ch))
	}
	return typ, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType1(prefix string) (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString(prefix)
	for tkn.lastChar != EOFCHAR {
		if tkn.lastChar == '\n' {
			tkn.ConsumeNext(buffer)
			break
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) scanCommentType2() (int, []byte) {
	buffer := bytes.NewBuffer(make([]byte, 0, 48))
	buffer.WriteString("/*")
	for {
		if tkn.lastChar == '*' {
			tkn.ConsumeNext(buffer)
			if tkn.lastChar == '/' {
				tkn.ConsumeNext(buffer)
				break
			}
			continue
		}
		if tkn.lastChar == EOFCHAR {
			return LEX_ERROR, buffer.Bytes()
		}
		tkn.ConsumeNext(buffer)
	}
	return COMMENT, buffer.Bytes()
}

func (tkn *Tokenizer) ConsumeNext(buffer *bytes.Buffer) {
	if tkn.lastChar == EOFCHAR {
		// This should never happen.
		panic("unexpected EOF")
	}
	buffer.WriteByte(byte(tkn.lastChar))
	tkn.next()
}

func (tkn *Tokenizer) next() {
	if ch, err := tkn.InStream.ReadByte(); err != nil {
		// Only EOF is possible.
		tkn.lastChar = EOFCHAR
	} else {
		tkn.lastChar = uint16(ch)
	}
	tkn.Position++
}

func isLetter(ch uint16) bool {
	return 'a' <= ch && ch <= 'z' || 'A' <= ch && ch <= 'Z' || ch == '_' || ch == '@'
}

func digitVal(ch uint16) int {
	switch {
	case '0' <= ch && ch <= '9':
		return int(ch) - '0'
	case 'a' <= ch && ch <= 'f':
		return int(ch) - 'a' + 10
	case 'A' <= ch && ch <= 'F':
		return int(ch) - 'A' + 10
	}
	return 16 // larger than any legal digit val
}

func isDigit(ch uint16) bool {
	return '0' <= ch && ch <= '9'
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io"
	"log"
	"os"

	"github.com/dustin/go-rs232"
)

func parseMode(s string) (rv rs232.SerConf) {
	switch s {
	case "8N1":
		rv = rs232.S_8N1
	case "7E1":
		rv = rs232.S_7E1
	case "7O1":
		rv = rs232.S_7O1
	default:
		log.Fatalf("Invalid mode: %v", s)
	}
	return
}

func main() {
	baudRate := flag.Int("baud", 57600, "Baud rate")
	mode := flag.String("mode", "8N1", "8N1 | 7E1 | 7O1")
	flag.Parse()

	port, err := rs232.OpenPort(flag.Arg(0), *baudRate,
		parseMode(*mode))
	if err != nil {
		log.Fatalf("Error opening port: %s", err)
	}

	io.Copy(os.Stdout, port)
}
<commit_msg>Require serialcat port argument and print better usage<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"log"
	"os"

	"github.com/dustin/go-rs232"
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [args] /dev/tty.whatever\n", os.Args[0])
		flag.PrintDefaults()
	}
}

func parseMode(s string) (rv rs232.SerConf) {
	switch s {
	case "8N1":
		rv = rs232.S_8N1
	case "7E1":
		rv = rs232.S_7E1
	case "7O1":
		rv = rs232.S_7O1
	default:
		log.Fatalf("Invalid mode: %v", s)
	}
	return
}

func main() {
	baudRate := flag.Int("baud", 57600, "Baud rate")
	mode := flag.String("mode", "8N1", "8N1 | 7E1 | 7O1")
	flag.Parse()

	path := flag.Arg(0)
	if path == "" {
		flag.Usage()
		os.Exit(64)
	}

	port, err := rs232.OpenPort(path, *baudRate, parseMode(*mode))
	if err != nil {
		log.Fatalf("Error opening port %q: %s", port, err)
	}

	io.Copy(os.Stdout, port)
}
<|endoftext|>


<commit_before>package whois
<commit_msg>Remove unused file.<commit_after><|endoftext|>


<commit_before>package registry

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
)

type tokenResponse struct {
	Token string `json:"token"`
}

func getToken(username, password string, params map[string]string, registryEndpoint *Endpoint) (token string, err error) {
	realm, ok := params["realm"]
	if !ok {
		return "", errors.New("no realm specified for token auth challenge")
	}

	realmURL, err := url.Parse(realm)
	if err != nil {
		return "", fmt.Errorf("invalid token auth challenge realm: %s", err)
	}

	if realmURL.Scheme == "" {
		if registryEndpoint.IsSecure {
			realmURL.Scheme = "https"
		} else {
			realmURL.Scheme = "http"
		}
	}

	req, err := http.NewRequest("GET", realmURL.String(), nil)
	if err != nil {
		return "", err
	}

	reqParams := req.URL.Query()
	service := params["service"]
	scope := params["scope"]

	if service != "" {
		reqParams.Add("service", service)
	}

	for _, scopeField := range strings.Fields(scope) {
		reqParams.Add("scope", scopeField)
	}

	if username != "" {
		reqParams.Add("account", username)
		req.SetBasicAuth(username, password)
	}

	req.URL.RawQuery = reqParams.Encode()

	resp, err := registryEndpoint.client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("token auth attempt for registry %s: %s request failed with status: %d %s", registryEndpoint, req.URL, resp.StatusCode, http.StatusText(resp.StatusCode))
	}

	decoder := json.NewDecoder(resp.Body)

	tr := new(tokenResponse)
	if err = decoder.Decode(tr); err != nil {
		return "", fmt.Errorf("unable to decode token response: %s", err)
	}

	if tr.Token == "" {
		return "", errors.New("authorization server did not include a token in the response")
	}

	return tr.Token, nil
}
<commit_msg>Registry: remove unwanted return variable name<commit_after>package registry

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"
)

type tokenResponse struct {
	Token string `json:"token"`
}

func getToken(username, password string, params map[string]string, registryEndpoint *Endpoint) (string, error) {
	realm, ok := params["realm"]
	if !ok {
		return "", errors.New("no realm specified for token auth challenge")
	}

	realmURL, err := url.Parse(realm)
	if err != nil {
		return "", fmt.Errorf("invalid token auth challenge realm: %s", err)
	}

	if realmURL.Scheme == "" {
		if registryEndpoint.IsSecure {
			realmURL.Scheme = "https"
		} else {
			realmURL.Scheme = "http"
		}
	}

	req, err := http.NewRequest("GET", realmURL.String(), nil)
	if err != nil {
		return "", err
	}

	reqParams := req.URL.Query()
	service := params["service"]
	scope := params["scope"]

	if service != "" {
		reqParams.Add("service", service)
	}

	for _, scopeField := range strings.Fields(scope) {
		reqParams.Add("scope", scopeField)
	}

	if username != "" {
		reqParams.Add("account", username)
		req.SetBasicAuth(username, password)
	}

	req.URL.RawQuery = reqParams.Encode()

	resp, err := registryEndpoint.client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("token auth attempt for registry %s: %s request failed with status: %d %s", registryEndpoint, req.URL, resp.StatusCode, http.StatusText(resp.StatusCode))
	}

	decoder := json.NewDecoder(resp.Body)

	tr := new(tokenResponse)
	if err = decoder.Decode(tr); err != nil {
		return "", fmt.Errorf("unable to decode token response: %s", err)
	}

	if tr.Token == "" {
		return "", errors.New("authorization server did not include a token in the response")
	}

	return tr.Token, nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2016 Yoshihiro Tanaka
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Author: Yoshihiro Tanaka <contact@cordea.jp>
 * date  : 2016-05-29
 */

package main

import (
    "flag"
    "log"
    "os"
    "io"
    "io/ioutil"
    "path/filepath"
    "strings"
)

var (
    isOverwrite = flag.Bool("w", false, "If True, overwrite the file if the file already exists in the target directory.")
)

const (
    DrawableDirPrefix = "drawable-"
)

func getFiles(path string) []os.FileInfo {
    stat, err := os.Stat(path)
    if err != nil {
        log.Fatalln(err)
    }
    switch md := stat.Mode(); {
    case md.IsDir():
        files, err := ioutil.ReadDir(path)
        if err != nil {
            log.Fatalln(err)
        }
        return files
    }
    return []os.FileInfo{}
}

func copyFile(src, trg string, isOverwrite bool) {
    if !isOverwrite {
        if _, err := os.Stat(trg); err == nil {
            log.Println(trg + " already exists. Not overwrite.")
            return
        }
    }

    srcFile, err := os.Open(src)
    if err != nil {
        log.Println(trg + ": " + err.Error())
    }
    defer srcFile.Close()

    trgFile, err := os.Create(trg)
    if err != nil {
        log.Println(trg + ": " + err.Error())
    }
    defer trgFile.Close()
    if _, err := io.Copy(trgFile, srcFile); err != nil {
        log.Println(trg + ": " + err.Error())
    }
    log.Println("Copied " + trg)
}

func isDrawableDir(dirname string) bool {
    return strings.Contains(dirname, DrawableDirPrefix)
}

func copyFiles(src, trg string, isOverwrite bool) {
    srcs := getFiles(src)
    trgs := getFiles(trg)

    if len(trgs) == 0 {
        for _, sf := range srcs {
            if isDrawableDir(sf.Name()) {
                os.Mkdir(filepath.Join(trg, sf.Name()), 0700)
            }
        }
        trgs = getFiles(trg)
    }

    for _, sf := range srcs {
        if !isDrawableDir(sf.Name()) {
            continue
        }
        for _, tf := range trgs {
            if sf.Name() == tf.Name() {
                sImages := getFiles(filepath.Join(src, sf.Name()))
                for _, si := range sImages {
                    sfPath := filepath.Join(src, sf.Name(), si.Name())
                    tfPath := filepath.Join(trg, tf.Name(), si.Name())
                    copyFile(sfPath, tfPath, isOverwrite)
                }
            }
        }
    }
}

func main() {
    flag.Parse()
    if flag.NArg() < 2 {
        log.Fatalln("Required argument is missing.")
    }

    src := flag.Arg(0)
    trg := flag.Arg(1)
    copyFiles(src, trg, *isOverwrite)
}
<commit_msg>add new option to drawable-copier<commit_after>/*
 * Copyright 2016 Yoshihiro Tanaka
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Author: Yoshihiro Tanaka <contact@cordea.jp>
 * date  : 2016-05-29
 */

package main

import (
    "flag"
    "log"
    "os"
    "io"
    "io/ioutil"
    "path/filepath"
    "strings"
)

var (
    isOverwrite = flag.Bool("w", false, "If True, overwrite the file if the file already exists in the target directory.")
    isRemoveSource = flag.Bool("r", false, "If True, remove the source directory when all of process was successful.")
)

const (
    DrawableDirPrefix = "drawable-"
)

func getFiles(path string) []os.FileInfo {
    stat, err := os.Stat(path)
    if err != nil {
        log.Fatalln(err)
    }
    switch md := stat.Mode(); {
    case md.IsDir():
        files, err := ioutil.ReadDir(path)
        if err != nil {
            log.Fatalln(err)
        }
        return files
    }
    return []os.FileInfo{}
}

func copyFile(src, trg string, isOverwrite bool) bool {
    if !isOverwrite {
        if _, err := os.Stat(trg); err == nil {
            log.Println(trg + " already exists. Not overwrite.")
            return true
        }
    }

    srcFile, err := os.Open(src)
    if err != nil {
        log.Println(trg + ": " + err.Error())
        return false
    }
    defer srcFile.Close()

    trgFile, err := os.Create(trg)
    if err != nil {
        log.Println(trg + ": " + err.Error())
        return false
    }
    defer trgFile.Close()
    if _, err := io.Copy(trgFile, srcFile); err != nil {
        os.Remove(trg)
        log.Println(trg + ": " + err.Error())
        return false
    }
    log.Println("Copied " + trg)
    return true
}

func isDrawableDir(dirname string) bool {
    return strings.Contains(dirname, DrawableDirPrefix)
}

func copyFiles(src, trg string, isOverwrite bool) bool {
    srcs := getFiles(src)
    trgs := getFiles(trg)

    if len(trgs) == 0 {
        for _, sf := range srcs {
            if isDrawableDir(sf.Name()) {
                os.Mkdir(filepath.Join(trg, sf.Name()), 0700)
            }
        }
        trgs = getFiles(trg)
    }

    allSucceed := true

    for _, sf := range srcs {
        if !isDrawableDir(sf.Name()) {
            continue
        }
        for _, tf := range trgs {
            if sf.Name() == tf.Name() {
                sImages := getFiles(filepath.Join(src, sf.Name()))
                for _, si := range sImages {
                    sfPath := filepath.Join(src, sf.Name(), si.Name())
                    tfPath := filepath.Join(trg, tf.Name(), si.Name())
                    allSucceed = copyFile(sfPath, tfPath, isOverwrite)
                }
            }
        }
    }
    return allSucceed
}

func main() {
    flag.Parse()
    if flag.NArg() < 2 {
        log.Fatalln("Required argument is missing.")
    }

    src := flag.Arg(0)
    trg := flag.Arg(1)
    allSucceed := copyFiles(src, trg, *isOverwrite)
    if *isRemoveSource && allSucceed {
        if err := os.RemoveAll(src); err != nil {
            log.Println(err)
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package vppcalls

import (
	"fmt"

	"time"

	"github.com/ligato/cn-infra/logging/measure"
	"github.com/ligato/vpp-agent/plugins/defaultplugins/ifplugin/bin_api/af_packet"
	intf "github.com/ligato/vpp-agent/plugins/defaultplugins/ifplugin/model/interfaces"
)

// AddAfPacketInterface calls AfPacketCreate VPP binary API.
func AddAfPacketInterface(afPacketIntf *intf.Interfaces_Interface_Afpacket, vppChan VPPChannel, timeLog measure.StopWatchEntry) (swIndex uint32, err error) {
	// AfPacketCreate time measurement
	start := time.Now()
	defer func() {
		if timeLog != nil {
			timeLog.LogTimeEntry(time.Since(start))
		}
	}()

	// Prepare the message.
	req := &af_packet.AfPacketCreate{
		HostIfName:      []byte(afPacketIntf.HostIfName),
		UseRandomHwAddr: 1,
	}

	reply := &af_packet.AfPacketCreateReply{}
	err = vppChan.SendRequest(req).ReceiveReply(reply)
	if err != nil {
		return 0, err
	}
	if reply.Retval != 0 {
		return 0, fmt.Errorf("add af_packet interface (%+v) returned %d", req, reply.Retval)
	}

	return reply.SwIfIndex, nil
}

// DeleteAfPacketInterface calls AfPacketDelete VPP binary API.
func DeleteAfPacketInterface(afPacketIntf *intf.Interfaces_Interface_Afpacket, vppChan VPPChannel, timeLog measure.StopWatchEntry) error {
	// AfPacketDelete time measurement
	start := time.Now()
	defer func() {
		if timeLog != nil {
			timeLog.LogTimeEntry(time.Since(start))
		}
	}()

	// Prepare the message.
	req := &af_packet.AfPacketDelete{
		HostIfName: []byte(afPacketIntf.HostIfName),
	}

	reply := &af_packet.AfPacketDeleteReply{}
	err := vppChan.SendRequest(req).ReceiveReply(reply)
	if err != nil {
		return err
	}
	if reply.Retval != 0 {
		return fmt.Errorf("deleting of af_packet interface (%+v) returned %d", req, reply.Retval)
	}

	return nil
}
<commit_msg>Fix error message info about vpp call request<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package vppcalls

import (
	"fmt"

	"time"

	"github.com/ligato/cn-infra/logging/measure"
	"github.com/ligato/vpp-agent/plugins/defaultplugins/ifplugin/bin_api/af_packet"
	intf "github.com/ligato/vpp-agent/plugins/defaultplugins/ifplugin/model/interfaces"
)

// AddAfPacketInterface calls AfPacketCreate VPP binary API.
func AddAfPacketInterface(afPacketIntf *intf.Interfaces_Interface_Afpacket, vppChan VPPChannel, timeLog measure.StopWatchEntry) (swIndex uint32, err error) {
	// AfPacketCreate time measurement
	start := time.Now()
	defer func() {
		if timeLog != nil {
			timeLog.LogTimeEntry(time.Since(start))
		}
	}()

	// Prepare the message.
	req := &af_packet.AfPacketCreate{
		HostIfName:      []byte(afPacketIntf.HostIfName),
		UseRandomHwAddr: 1,
	}

	reply := &af_packet.AfPacketCreateReply{}
	err = vppChan.SendRequest(req).ReceiveReply(reply)
	if err != nil {
		return 0, err
	}
	if reply.Retval != 0 {
		return 0, fmt.Errorf("add af_packet interface (%+v) returned %d", afPacketIntf, reply.Retval)
	}

	return reply.SwIfIndex, nil
}

// DeleteAfPacketInterface calls AfPacketDelete VPP binary API.
func DeleteAfPacketInterface(afPacketIntf *intf.Interfaces_Interface_Afpacket, vppChan VPPChannel, timeLog measure.StopWatchEntry) error {
	// AfPacketDelete time measurement
	start := time.Now()
	defer func() {
		if timeLog != nil {
			timeLog.LogTimeEntry(time.Since(start))
		}
	}()

	// Prepare the message.
	req := &af_packet.AfPacketDelete{
		HostIfName: []byte(afPacketIntf.HostIfName),
	}

	reply := &af_packet.AfPacketDeleteReply{}
	err := vppChan.SendRequest(req).ReceiveReply(reply)
	if err != nil {
		return err
	}
	if reply.Retval != 0 {
		return fmt.Errorf("deleting of af_packet interface (%+v) returned %d", afPacketIntf, reply.Retval)
	}

	return nil
}
<|endoftext|>


<commit_before>package shell

import (
	"context"
	"fmt"
	"io"
	"math/rand"
	"sort"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
)

func init() {
	Commands = append(Commands, &commandVolumeFixReplication{})
}

type commandVolumeFixReplication struct {
}

func (c *commandVolumeFixReplication) Name() string {
	return "volume.fix.replication"
}

func (c *commandVolumeFixReplication) Help() string {
	return `add replicas to volumes that are missing replicas

	This command finds all under-replicated volumes, and finds volume servers with free slots.
	If the free slots satisfy the replication requirement, the volume content is copied over and mounted.

	volume.fix.replication -n # do not take action
	volume.fix.replication    # actually copying the volume files and mount the volume

	Note:
		* each time this will only add back one replica for one volume id. If there are multiple replicas
		  are missing, e.g. multiple volume servers are new, you may need to run this multiple times.
		* do not run this too quickly within seconds, since the new volume replica may take a few seconds
		  to register itself to the master.

`
}

func (c *commandVolumeFixReplication) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	takeAction := true
	if len(args) > 0 && args[0] == "-n" {
		takeAction = false
	}

	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return err
	}

	// find all volumes that needs replication
	// collect all data nodes
	replicatedVolumeLocations := make(map[uint32][]location)
	replicatedVolumeInfo := make(map[uint32]*master_pb.VolumeInformationMessage)
	var allLocations []location
	eachDataNode(resp.TopologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		loc := newLocation(dc, string(rack), dn)
		for _, v := range dn.VolumeInfos {
			if v.ReplicaPlacement > 0 {
				replicatedVolumeLocations[v.Id] = append(replicatedVolumeLocations[v.Id], loc)
				replicatedVolumeInfo[v.Id] = v
			}
		}
		allLocations = append(allLocations, loc)
	})

	// find all under replicated volumes
	underReplicatedVolumeLocations := make(map[uint32][]location)
	overReplicatedVolumeLocations := make(map[uint32][]location)
	for vid, locations := range replicatedVolumeLocations {
		volumeInfo := replicatedVolumeInfo[vid]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(volumeInfo.ReplicaPlacement))
		if replicaPlacement.GetCopyCount() > len(locations) {
			underReplicatedVolumeLocations[vid] = locations
		} else if replicaPlacement.GetCopyCount() < len(locations) {
			overReplicatedVolumeLocations[vid] = locations
			fmt.Fprintf(writer, "volume %d replication %s, but over replicated:%+v\n", volumeInfo.Id, replicaPlacement, locations)
		}
	}

	if len(underReplicatedVolumeLocations) == 0 {
		return fmt.Errorf("no under replicated volumes")
	}

	if len(allLocations) == 0 {
		return fmt.Errorf("no data nodes at all")
	}

	// find the most under populated data nodes
	keepDataNodesSorted(allLocations)

	return c.fixUnderReplicatedVolumes(commandEnv, writer, takeAction, underReplicatedVolumeLocations, replicatedVolumeInfo, allLocations)

}

func (c *commandVolumeFixReplication) fixUnderReplicatedVolumes(commandEnv *CommandEnv, writer io.Writer, takeAction bool, underReplicatedVolumeLocations map[uint32][]location, replicatedVolumeInfo map[uint32]*master_pb.VolumeInformationMessage, allLocations []location) error {
	for vid, locations := range underReplicatedVolumeLocations {
		volumeInfo := replicatedVolumeInfo[vid]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(volumeInfo.ReplicaPlacement))
		foundNewLocation := false
		for _, dst := range allLocations {
			// check whether data nodes satisfy the constraints
			if dst.dataNode.FreeVolumeCount > 0 && satisfyReplicaPlacement(replicaPlacement, locations, dst) {
				// ask the volume server to replicate the volume
				sourceNodes := underReplicatedVolumeLocations[vid]
				sourceNode := sourceNodes[rand.Intn(len(sourceNodes))]
				foundNewLocation = true
				fmt.Fprintf(writer, "replicating volume %d %s from %s to dataNode %s ...\n", volumeInfo.Id, replicaPlacement, sourceNode.dataNode.Id, dst.dataNode.Id)

				if !takeAction {
					break
				}

				err := operation.WithVolumeServerClient(dst.dataNode.Id, commandEnv.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
					_, replicateErr := volumeServerClient.VolumeCopy(context.Background(), &volume_server_pb.VolumeCopyRequest{
						VolumeId:       volumeInfo.Id,
						SourceDataNode: sourceNode.dataNode.Id,
					})
					if replicateErr != nil {
						return fmt.Errorf("copying from %s => %s : %v", sourceNode.dataNode.Id, dst.dataNode.Id, replicateErr)
					}
					return nil
				})

				if err != nil {
					return err
				}

				// adjust free volume count
				dst.dataNode.FreeVolumeCount--
				keepDataNodesSorted(allLocations)
				break
			}
		}
		if !foundNewLocation {
			fmt.Fprintf(writer, "failed to place volume %d replica as %s, existing:%+v\n", volumeInfo.Id, replicaPlacement, locations)
		}

	}
	return nil
}

func keepDataNodesSorted(dataNodes []location) {
	sort.Slice(dataNodes, func(i, j int) bool {
		return dataNodes[i].dataNode.FreeVolumeCount > dataNodes[j].dataNode.FreeVolumeCount
	})
}

/*
  if on an existing data node {
    return false
  }
  if different from existing dcs {
    if lack on different dcs {
      return true
    }else{
      return false
    }
  }
  if not on primary dc {
    return false
  }
  if different from existing racks {
    if lack on different racks {
      return true
    }else{
      return false
    }
  }
  if not on primary rack {
    return false
  }
  if lacks on same rack {
    return true
  } else {
    return false
  }
*/
func satisfyReplicaPlacement(replicaPlacement *super_block.ReplicaPlacement, existingLocations []location, possibleLocation location) bool {

	existingDataNodes := make(map[string]int)
	for _, loc := range existingLocations {
		existingDataNodes[loc.String()] += 1
	}
	sameDataNodeCount := existingDataNodes[possibleLocation.String()]
	// avoid duplicated volume on the same data node
	if sameDataNodeCount > 0 {
		return false
	}

	existingDataCenters := make(map[string]int)
	for _, loc := range existingLocations {
		existingDataCenters[loc.DataCenter()] += 1
	}
	primaryDataCenters, _ := findTopKeys(existingDataCenters)

	// ensure data center count is within limit
	if _, found := existingDataCenters[possibleLocation.DataCenter()]; !found {
		// different from existing dcs
		if len(existingDataCenters) < replicaPlacement.DiffDataCenterCount+1 {
			// lack on different dcs
			return true
		} else {
			// adding this would go over the different dcs limit
			return false
		}
	}
	// now this is same as one of the existing data center
	if !isAmong(possibleLocation.DataCenter(), primaryDataCenters) {
		// not on one of the primary dcs
		return false
	}

	// now this is one of the primary dcs
	existingRacks := make(map[string]int)
	for _, loc := range existingLocations {
		if loc.DataCenter() != possibleLocation.DataCenter() {
			continue
		}
		existingRacks[loc.Rack()] += 1
	}
	primaryRacks, _ := findTopKeys(existingRacks)
	sameRackCount := existingRacks[possibleLocation.Rack()]

	// ensure rack count is within limit
	if _, found := existingRacks[possibleLocation.Rack()]; !found {
		// different from existing racks
		if len(existingRacks) < replicaPlacement.DiffRackCount+1 {
			// lack on different racks
			return true
		} else {
			// adding this would go over the different racks limit
			return false
		}
	}
	// now this is same as one of the existing racks
	if !isAmong(possibleLocation.Rack(), primaryRacks) {
		// not on the primary rack
		return false
	}

	// now this is on the primary rack

	// different from existing data nodes
	if sameRackCount < replicaPlacement.SameRackCount+1 {
		// lack on same rack
		return true
	} else {
		// adding this would go over the same data node limit
		return false
	}

}

func findTopKeys(m map[string]int) (topKeys []string, max int) {
	for k, c := range m {
		if max < c {
			topKeys = topKeys[:0]
			topKeys = append(topKeys, k)
			max = c
		} else if max == c {
			topKeys = append(topKeys, k)
		}
	}
	return
}

func isAmong(key string, keys []string) bool {
	for _, k := range keys {
		if k == key {
			return true
		}
	}
	return false
}

type location struct {
	dc       string
	rack     string
	dataNode *master_pb.DataNodeInfo
}

func newLocation(dc, rack string, dataNode *master_pb.DataNodeInfo) location {
	return location{
		dc:       dc,
		rack:     rack,
		dataNode: dataNode,
	}
}

func (l location) String() string {
	return fmt.Sprintf("%s %s %s", l.dc, l.rack, l.dataNode.Id)
}

func (l location) Rack() string {
	return fmt.Sprintf("%s %s", l.dc, l.rack)
}

func (l location) DataCenter() string {
	return l.dc
}
<commit_msg>refactoring<commit_after>package shell

import (
	"context"
	"fmt"
	"io"
	"math/rand"
	"sort"

	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/super_block"
)

func init() {
	Commands = append(Commands, &commandVolumeFixReplication{})
}

type commandVolumeFixReplication struct {
}

func (c *commandVolumeFixReplication) Name() string {
	return "volume.fix.replication"
}

func (c *commandVolumeFixReplication) Help() string {
	return `add replicas to volumes that are missing replicas

	This command finds all under-replicated volumes, and finds volume servers with free slots.
	If the free slots satisfy the replication requirement, the volume content is copied over and mounted.

	volume.fix.replication -n # do not take action
	volume.fix.replication    # actually copying the volume files and mount the volume

	Note:
		* each time this will only add back one replica for one volume id. If there are multiple replicas
		  are missing, e.g. multiple volume servers are new, you may need to run this multiple times.
		* do not run this too quickly within seconds, since the new volume replica may take a few seconds
		  to register itself to the master.

`
}

func (c *commandVolumeFixReplication) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	if err = commandEnv.confirmIsLocked(); err != nil {
		return
	}

	takeAction := true
	if len(args) > 0 && args[0] == "-n" {
		takeAction = false
	}

	var resp *master_pb.VolumeListResponse
	err = commandEnv.MasterClient.WithClient(func(client master_pb.SeaweedClient) error {
		resp, err = client.VolumeList(context.Background(), &master_pb.VolumeListRequest{})
		return err
	})
	if err != nil {
		return err
	}

	// find all volumes that needs replication
	// collect all data nodes
	volumeReplicas := make(map[uint32][]*VolumeReplica)
	var allLocations []location
	eachDataNode(resp.TopologyInfo, func(dc string, rack RackId, dn *master_pb.DataNodeInfo) {
		loc := newLocation(dc, string(rack), dn)
		for _, v := range dn.VolumeInfos {
			if v.ReplicaPlacement > 0 {
				volumeReplicas[v.Id] = append(volumeReplicas[v.Id], &VolumeReplica{
					location: &loc,
					info:     v,
				})
			}
		}
		allLocations = append(allLocations, loc)
	})

	// find all under replicated volumes
	var underReplicatedVolumeIds, overReplicatedVolumeIds []uint32
	for vid, replicas := range volumeReplicas {
		replica := replicas[rand.Intn(len(replicas))]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replica.info.ReplicaPlacement))
		if replicaPlacement.GetCopyCount() > len(replicas) {
			underReplicatedVolumeIds = append(underReplicatedVolumeIds, vid)
		} else if replicaPlacement.GetCopyCount() < len(replicas) {
			overReplicatedVolumeIds = append(overReplicatedVolumeIds, vid)
			fmt.Fprintf(writer, "volume %d replication %s, but over replicated %+d\n", replica.info.Id, replicaPlacement, len(replicas))
		}
	}

	if len(underReplicatedVolumeIds) == 0 {
		return fmt.Errorf("no under replicated volumes")
	}

	if len(allLocations) == 0 {
		return fmt.Errorf("no data nodes at all")
	}

	// find the most under populated data nodes
	keepDataNodesSorted(allLocations)

	return c.fixUnderReplicatedVolumes(commandEnv, writer, takeAction, underReplicatedVolumeIds, volumeReplicas, allLocations)

}

func (c *commandVolumeFixReplication) fixUnderReplicatedVolumes(commandEnv *CommandEnv, writer io.Writer, takeAction bool, underReplicatedVolumeIds []uint32, volumeReplicas map[uint32][]*VolumeReplica, allLocations []location) error {
	for _, vid := range underReplicatedVolumeIds {
		replicas := volumeReplicas[vid]
		replica := replicas[rand.Intn(len(replicas))]
		replicaPlacement, _ := super_block.NewReplicaPlacementFromByte(byte(replica.info.ReplicaPlacement))
		foundNewLocation := false
		for _, dst := range allLocations {
			// check whether data nodes satisfy the constraints
			if dst.dataNode.FreeVolumeCount > 0 && satisfyReplicaPlacement(replicaPlacement, replicas, dst) {
				// ask the volume server to replicate the volume
				foundNewLocation = true
				fmt.Fprintf(writer, "replicating volume %d %s from %s to dataNode %s ...\n", replica.info.Id, replicaPlacement, replica.location.dataNode.Id, dst.dataNode.Id)

				if !takeAction {
					break
				}

				err := operation.WithVolumeServerClient(dst.dataNode.Id, commandEnv.option.GrpcDialOption, func(volumeServerClient volume_server_pb.VolumeServerClient) error {
					_, replicateErr := volumeServerClient.VolumeCopy(context.Background(), &volume_server_pb.VolumeCopyRequest{
						VolumeId:       replica.info.Id,
						SourceDataNode: replica.location.dataNode.Id,
					})
					if replicateErr != nil {
						return fmt.Errorf("copying from %s => %s : %v", replica.location.dataNode.Id, dst.dataNode.Id, replicateErr)
					}
					return nil
				})

				if err != nil {
					return err
				}

				// adjust free volume count
				dst.dataNode.FreeVolumeCount--
				keepDataNodesSorted(allLocations)
				break
			}
		}
		if !foundNewLocation {
			fmt.Fprintf(writer, "failed to place volume %d replica as %s, existing:%+v\n", replica.info.Id, replicaPlacement, len(replicas))
		}

	}
	return nil
}

func keepDataNodesSorted(dataNodes []location) {
	sort.Slice(dataNodes, func(i, j int) bool {
		return dataNodes[i].dataNode.FreeVolumeCount > dataNodes[j].dataNode.FreeVolumeCount
	})
}

/*
  if on an existing data node {
    return false
  }
  if different from existing dcs {
    if lack on different dcs {
      return true
    }else{
      return false
    }
  }
  if not on primary dc {
    return false
  }
  if different from existing racks {
    if lack on different racks {
      return true
    }else{
      return false
    }
  }
  if not on primary rack {
    return false
  }
  if lacks on same rack {
    return true
  } else {
    return false
  }
*/
func satisfyReplicaPlacement(replicaPlacement *super_block.ReplicaPlacement, replicas []*VolumeReplica, possibleLocation location) bool {

	existingDataNodes := make(map[string]int)
	for _, replica := range replicas {
		existingDataNodes[replica.location.String()] += 1
	}
	sameDataNodeCount := existingDataNodes[possibleLocation.String()]
	// avoid duplicated volume on the same data node
	if sameDataNodeCount > 0 {
		return false
	}

	existingDataCenters := make(map[string]int)
	for _, replica := range replicas {
		existingDataCenters[replica.location.DataCenter()] += 1
	}
	primaryDataCenters, _ := findTopKeys(existingDataCenters)

	// ensure data center count is within limit
	if _, found := existingDataCenters[possibleLocation.DataCenter()]; !found {
		// different from existing dcs
		if len(existingDataCenters) < replicaPlacement.DiffDataCenterCount+1 {
			// lack on different dcs
			return true
		} else {
			// adding this would go over the different dcs limit
			return false
		}
	}
	// now this is same as one of the existing data center
	if !isAmong(possibleLocation.DataCenter(), primaryDataCenters) {
		// not on one of the primary dcs
		return false
	}

	// now this is one of the primary dcs
	existingRacks := make(map[string]int)
	for _, replica := range replicas {
		if replica.location.DataCenter() != possibleLocation.DataCenter() {
			continue
		}
		existingRacks[replica.location.Rack()] += 1
	}
	primaryRacks, _ := findTopKeys(existingRacks)
	sameRackCount := existingRacks[possibleLocation.Rack()]

	// ensure rack count is within limit
	if _, found := existingRacks[possibleLocation.Rack()]; !found {
		// different from existing racks
		if len(existingRacks) < replicaPlacement.DiffRackCount+1 {
			// lack on different racks
			return true
		} else {
			// adding this would go over the different racks limit
			return false
		}
	}
	// now this is same as one of the existing racks
	if !isAmong(possibleLocation.Rack(), primaryRacks) {
		// not on the primary rack
		return false
	}

	// now this is on the primary rack

	// different from existing data nodes
	if sameRackCount < replicaPlacement.SameRackCount+1 {
		// lack on same rack
		return true
	} else {
		// adding this would go over the same data node limit
		return false
	}

}

func findTopKeys(m map[string]int) (topKeys []string, max int) {
	for k, c := range m {
		if max < c {
			topKeys = topKeys[:0]
			topKeys = append(topKeys, k)
			max = c
		} else if max == c {
			topKeys = append(topKeys, k)
		}
	}
	return
}

func isAmong(key string, keys []string) bool {
	for _, k := range keys {
		if k == key {
			return true
		}
	}
	return false
}

type VolumeReplica struct {
	location *location
	info     *master_pb.VolumeInformationMessage
}

type location struct {
	dc       string
	rack     string
	dataNode *master_pb.DataNodeInfo
}

func newLocation(dc, rack string, dataNode *master_pb.DataNodeInfo) location {
	return location{
		dc:       dc,
		rack:     rack,
		dataNode: dataNode,
	}
}

func (l location) String() string {
	return fmt.Sprintf("%s %s %s", l.dc, l.rack, l.dataNode.Id)
}

func (l location) Rack() string {
	return fmt.Sprintf("%s %s", l.dc, l.rack)
}

func (l location) DataCenter() string {
	return l.dc
}
<|endoftext|>