<commit_before>package jsonrpc2_test

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/rpc"
	"reflect"
	"strings"
	"testing"

	"github.com/powerman/rpc-codec/jsonrpc2"
)

// Svc is an RPC service for testing.
type Svc struct{}

func (*Svc) Sum(vals [2]int, res *int) error {
	*res = vals[0] + vals[1]
	return nil
}

func init() {
	_ = rpc.Register(&Svc{})
}

func TestHTTPServer(t *testing.T) {
	const jBad = `{}`
	const jSum = `{"jsonrpc":"2.0","id":0,"method":"Svc.Sum","params":[3,5]}`
	const jNotify = `{"jsonrpc":"2.0","method":"Svc.Sum","params":[3,5]}`
	const jRes = `{"jsonrpc":"2.0","id":0,"result":8}`
	const jErr = `{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"invalid request"}}`
	const jParse = `{"jsonrpc":"2.0","id":null,"error":{"code":-32700,"message":"parse error"}}`
	const contentType = "application/json"

	cases := []struct {
		method      string
		contentType string
		accept      string
		body        string
		code        int
		reply       string
	}{
		{"GET", "", "", "", http.StatusMethodNotAllowed, ""},
		{"POST", contentType, "", jSum, http.StatusUnsupportedMediaType, ""},
		{"POST", "text/json", contentType, jSum, http.StatusUnsupportedMediaType, ""},
		{"PUT", contentType, contentType, jSum, http.StatusMethodNotAllowed, ""},
		{"POST", contentType, contentType, jNotify, http.StatusNoContent, ""},
		{"POST", contentType, contentType, jSum, http.StatusOK, jRes},
		{"POST", contentType, contentType, jBad, http.StatusOK, jErr},
		{"POST", contentType, contentType, "", http.StatusOK, jParse},
		{"POST", contentType, contentType, " ", http.StatusOK, jParse},
		{"POST", contentType, contentType, "{", http.StatusOK, jParse},
		{"POST", contentType, contentType, `{"jsonrpc":"2.0",`, http.StatusOK, jParse},
		{"POST", contentType + "; charset=utf-8", contentType, jSum, http.StatusOK, jRes},
	}

	ts := httptest.NewServer(jsonrpc2.HTTPHandler(nil))
	// Don't close because of https://github.com/golang/go/issues/12262
	// defer ts.Close()

	for _, c := range cases {
		req, err := http.NewRequest(c.method, ts.URL, strings.NewReader(c.body))
		if err != nil {
			t.Errorf("NewRequest(%s %s), err = %v", c.method, ts.URL, err)
		}
		if c.contentType != "" {
			req.Header.Add("Content-Type", c.contentType)
		}
		if c.accept != "" {
			req.Header.Add("Accept", c.accept)
		}
		resp, err := (&http.Client{}).Do(req)
		if err != nil {
			t.Errorf("Do(%s %s), err = %v", c.method, ts.URL, err)
		}
		if resp.StatusCode != c.code {
			t.Errorf("Do(%s %s), status = %v, want = %v", c.method, ts.URL, resp.StatusCode, c.code)
		}
		if resp.Header.Get("Content-Type") != contentType {
			t.Errorf("Do(%s %s), Content-Type = %q, want = %q", c.method, ts.URL, resp.Header.Get("Content-Type"), contentType)
		}
		got, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Errorf("ReadAll(), err = %v", err)
		}
		if c.reply == "" {
			if len(got) != 0 {
				t.Errorf("Do(%s %s)\nexp: %#q\ngot: %#q", c.method, ts.URL, c.reply, string(bytes.TrimRight(got, "\n")))
			}
		} else {
			var jgot, jwant interface{}
			if err := json.Unmarshal(got, &jgot); err != nil {
				t.Errorf("Do(%s %s), output err = %v\ngot: %#q", c.method, ts.URL, err, string(bytes.TrimRight(got, "\n")))
			}
			if err := json.Unmarshal([]byte(c.reply), &jwant); err != nil {
				t.Errorf("Do(%s %s), expect err = %v\nexp: %#q", c.method, ts.URL, err, c.reply)
			}
			if !reflect.DeepEqual(jgot, jwant) {
				t.Errorf("Do(%s %s)\nexp: %#q\ngot: %#q", c.method, ts.URL, c.reply, string(bytes.TrimRight(got, "\n")))
			}
		}
	}
}

func TestHTTPClient(t *testing.T) {
	ts := httptest.NewServer(jsonrpc2.HTTPHandler(nil))
	// Don't close because of https://github.com/golang/go/issues/12262
	// defer ts.Close()
	client := jsonrpc2.NewHTTPClient(ts.URL)
	defer client.Close()

	var in [2]int
	var got, want int

	in, want = [2]int{1, 2}, 3
	err := client.Call("Svc.Sum", in, &got)
	if err != nil {
		t.Errorf("Call(%v), err = %v", in, err)
	}
	if got != want {
		t.Errorf("Call(%v) = %v, want = %v", in, got, want)
	}

	in = [2]int{2, 3}
	err = client.Notify("Svc.Sum", in)
	if err != nil {
		t.Errorf("Notify(%v), err = %v", in, err)
	}

	in, want = [2]int{3, 4}, 7
	err = client.Call("Svc.Sum", in, &got)
	if err != nil {
		t.Errorf("Call(%v), err = %v", in, err)
	}
	if got != want {
		t.Errorf("Call(%v) = %v, want = %v", in, got, want)
	}
}

type ContentTypeHandler string

func (h ContentTypeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", string(h))
	io.WriteString(w, `{"jsonrpc":"2.0","id":0,"result":8}`)
}

func TestHTTPClientContentType(t *testing.T) {
	const contentType = "application/json"
	cases := []struct {
		contentType string
		errorString string
	}{
		{contentType, "<nil>"},
		{contentType + "; charset=utf-8", "<nil>"},
		{contentType + "; bad", `{"code":-32603,"message":"bad HTTP Content-Type: application/json; bad"}`},
		{contentType + "fail", `{"code":-32603,"message":"bad HTTP Content-Type: application/jsonfail"}`},
	}

	for _, c := range cases {
		ts := httptest.NewServer(ContentTypeHandler(c.contentType))

		client := jsonrpc2.NewHTTPClient(ts.URL)
		in := []string{"ads"}
		var got int
		err := client.Call("Svc.Sum", in, &got)
		actualErrString := fmt.Sprintf("%v", err)
		if actualErrString != c.errorString {
			t.Errorf("Unexpected result. exp: %#q\ngot: %#q", actualErrString, c.errorString)
		}
	}
}
<commit_msg>improve test<commit_after>package jsonrpc2_test

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/rpc"
	"reflect"
	"strings"
	"testing"

	"github.com/powerman/rpc-codec/jsonrpc2"
)

// Svc is an RPC service for testing.
type Svc struct{}

func (*Svc) Sum(vals [2]int, res *int) error {
	*res = vals[0] + vals[1]
	return nil
}

func init() {
	_ = rpc.Register(&Svc{})
}

func TestHTTPServer(t *testing.T) {
	const jBad = `{}`
	const jSum = `{"jsonrpc":"2.0","id":0,"method":"Svc.Sum","params":[3,5]}`
	const jNotify = `{"jsonrpc":"2.0","method":"Svc.Sum","params":[3,5]}`
	const jRes = `{"jsonrpc":"2.0","id":0,"result":8}`
	const jErr = `{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"invalid request"}}`
	const jParse = `{"jsonrpc":"2.0","id":null,"error":{"code":-32700,"message":"parse error"}}`
	const contentType = "application/json"

	cases := []struct {
		method      string
		contentType string
		accept      string
		body        string
		code        int
		reply       string
	}{
		{"GET", "", "", "", http.StatusMethodNotAllowed, ""},
		{"POST", contentType, "", jSum, http.StatusUnsupportedMediaType, ""},
		{"POST", "text/json", contentType, jSum, http.StatusUnsupportedMediaType, ""},
		{"PUT", contentType, contentType, jSum, http.StatusMethodNotAllowed, ""},
		{"POST", contentType, contentType, jNotify, http.StatusNoContent, ""},
		{"POST", contentType, contentType, jSum, http.StatusOK, jRes},
		{"POST", contentType, contentType, jBad, http.StatusOK, jErr},
		{"POST", contentType, contentType, "", http.StatusOK, jParse},
		{"POST", contentType, contentType, " ", http.StatusOK, jParse},
		{"POST", contentType, contentType, "{", http.StatusOK, jParse},
		{"POST", contentType, contentType, `{"jsonrpc":"2.0",`, http.StatusOK, jParse},
		{"POST", contentType + "; charset=utf-8", contentType, jSum, http.StatusOK, jRes},
	}

	ts := httptest.NewServer(jsonrpc2.HTTPHandler(nil))
	// Don't close because of https://github.com/golang/go/issues/12262
	// defer ts.Close()

	for _, c := range cases {
		req, err := http.NewRequest(c.method, ts.URL, strings.NewReader(c.body))
		if err != nil {
			t.Errorf("NewRequest(%s %s), err = %v", c.method, ts.URL, err)
		}
		if c.contentType != "" {
			req.Header.Add("Content-Type", c.contentType)
		}
		if c.accept != "" {
			req.Header.Add("Accept", c.accept)
		}
		resp, err := (&http.Client{}).Do(req)
		if err != nil {
			t.Errorf("Do(%s %s), err = %v", c.method, ts.URL, err)
		}
		if resp.StatusCode != c.code {
			t.Errorf("Do(%s %s), status = %v, want = %v", c.method, ts.URL, resp.StatusCode, c.code)
		}
		if resp.Header.Get("Content-Type") != contentType {
			t.Errorf("Do(%s %s), Content-Type = %q, want = %q", c.method, ts.URL, resp.Header.Get("Content-Type"), contentType)
		}
		got, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			t.Errorf("ReadAll(), err = %v", err)
		}
		if c.reply == "" {
			if len(got) != 0 {
				t.Errorf("Do(%s %s)\nexp: %#q\ngot: %#q", c.method, ts.URL, c.reply, string(bytes.TrimRight(got, "\n")))
			}
		} else {
			var jgot, jwant interface{}
			if err := json.Unmarshal(got, &jgot); err != nil {
				t.Errorf("Do(%s %s), output err = %v\ngot: %#q", c.method, ts.URL, err, string(bytes.TrimRight(got, "\n")))
			}
			if err := json.Unmarshal([]byte(c.reply), &jwant); err != nil {
				t.Errorf("Do(%s %s), expect err = %v\nexp: %#q", c.method, ts.URL, err, c.reply)
			}
			if !reflect.DeepEqual(jgot, jwant) {
				t.Errorf("Do(%s %s)\nexp: %#q\ngot: %#q", c.method, ts.URL, c.reply, string(bytes.TrimRight(got, "\n")))
			}
		}
	}
}

func TestHTTPClient(t *testing.T) {
	ts := httptest.NewServer(jsonrpc2.HTTPHandler(nil))
	// Don't close because of https://github.com/golang/go/issues/12262
	// defer ts.Close()
	client := jsonrpc2.NewHTTPClient(ts.URL)
	defer client.Close()

	var in [2]int
	var got, want int

	in, want = [2]int{1, 2}, 3
	err := client.Call("Svc.Sum", in, &got)
	if err != nil {
		t.Errorf("Call(%v), err = %v", in, err)
	}
	if got != want {
		t.Errorf("Call(%v) = %v, want = %v", in, got, want)
	}

	in = [2]int{2, 3}
	err = client.Notify("Svc.Sum", in)
	if err != nil {
		t.Errorf("Notify(%v), err = %v", in, err)
	}

	in, want = [2]int{3, 4}, 7
	err = client.Call("Svc.Sum", in, &got)
	if err != nil {
		t.Errorf("Call(%v), err = %v", in, err)
	}
	if got != want {
		t.Errorf("Call(%v) = %v, want = %v", in, got, want)
	}
}

type ContentTypeHandler string

func (h ContentTypeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", string(h))
	io.WriteString(w, `{"jsonrpc":"2.0","id":0,"result":8}`)
}

func TestHTTPClientContentType(t *testing.T) {
	const contentType = "application/json"
	cases := []struct {
		contentType string
		wanterr     *jsonrpc2.Error
	}{
		{contentType, nil},
		{contentType + "; charset=utf-8", nil},
		{contentType + "; bad", jsonrpc2.NewError(-32603, "bad HTTP Content-Type: "+contentType+"; bad")},
		{contentType + "rpc", jsonrpc2.NewError(-32603, "bad HTTP Content-Type: "+contentType+"rpc")},
	}

	for _, v := range cases {
		ts := httptest.NewServer(ContentTypeHandler(v.contentType))
		client := jsonrpc2.NewHTTPClient(ts.URL)

		err := jsonrpc2.ServerError(client.Call("Svc.Sum", [2]int{3, 4}, nil))
		if fmt.Sprintf("%#v", err) != fmt.Sprintf("%#v", v.wanterr) {
			t.Errorf("%q, err = %#v, want %#v", v.contentType, err, v.wanterr)
		}

		client.Close()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"path"
	"strings"
	"time"

	"github.com/coreos/etcd/etcdserver"
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"github.com/coreos/etcd/mvcc/backend"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"github.com/coreos/etcd/pkg/pbutil"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/coreos/etcd/snap"
	"github.com/coreos/etcd/store"
	"github.com/coreos/etcd/wal"
	"github.com/coreos/etcd/wal/walpb"
)

var (
	migrateDatadir string
	ttl            time.Duration
)

func init() {
	flag.StringVar(&migrateDatadir, "data-dir", "", "Path to the data directory")
	flag.DurationVar(&ttl, "ttl", time.Hour, "TTL of event keys (default 1 hour)")
	flag.Parse()
}

func main() {
	dbpath := path.Join(migrateDatadir, "member", "snap", "db")

	be := backend.New(dbpath, time.Second, 10000)
	tx := be.BatchTx()

	st := store.New("/0", "/1")
	expireTime := time.Now().Add(ttl)

	tx.Lock()
	err := tx.UnsafeForEach([]byte("key"), func(k, v []byte) error {
		tb := isTombstone(k)

		kv := &mvccpb.KeyValue{}
		kv.Unmarshal(v)

		// This is a compact key..
		if !strings.HasPrefix(string(kv.Key), "/") {
			return nil
		}

		// fmt.Printf("%s %d %d %d\n", string(kv.Key), kv.CreateRevision, kv.ModRevision, kv.Lease)
		ttlOpt := store.TTLOptionSet{}
		if kv.Lease != 0 {
			ttlOpt = store.TTLOptionSet{ExpireTime: expireTime}
		}

		if !tb {
			_, err := st.Set(path.Join("1", string(kv.Key)), false, string(kv.Value), ttlOpt)
			if err != nil {
				return err
			}
		} else {
			st.Delete(string(kv.Key), false, false)
		}

		return nil
	})
	if err != nil {
		panic(err)
	}
	tx.Unlock()

	traverse(st, "/")

	snapshotter := snap.New(path.Join(migrateDatadir, "member", "snap"))
	raftSnap, err := snapshotter.Load()
	if err != nil {
		panic(err)
	}

	metadata, hardstate, oldSt := rebuild(migrateDatadir)

	if err := os.RemoveAll(migrateDatadir); err != nil {
		panic(err)
	}
	if err := os.MkdirAll(path.Join(migrateDatadir, "member", "snap"), 0700); err != nil {
		panic(err)
	}
	walDir := path.Join(migrateDatadir, "member", "wal")

	w, err := wal.Create(walDir, metadata)
	if err != nil {
		panic(err)
	}
	err = w.SaveSnapshot(walpb.Snapshot{Index: hardstate.Commit, Term: hardstate.Term})
	if err != nil {
		panic(err)
	}
	w.Close()

	event, err := oldSt.Get("/0", true, false)
	if err != nil {
		panic(err)
	}
	q := []*store.NodeExtern{}
	q = append(q, event.Node)
	for len(q) > 0 {
		n := q[0]
		q = q[1:]
		v := ""
		if !n.Dir {
			v = *n.Value
		}
		if n.Key != "/0" {
			if n.Key == path.Join("/0", "version") {
				v = "2.3.7"
			}
			if _, err := st.Set(n.Key, n.Dir, v, store.TTLOptionSet{}); err != nil {
				panic(err)
			}
		}
		for _, next := range n.Nodes {
			q = append(q, next)
		}
	}

	data, err := st.Save()
	if err != nil {
		panic(err)
	}
	raftSnap.Data = data
	if err := snapshotter.SaveSnap(*raftSnap); err != nil {
		panic(err)
	}
	fmt.Println("Finished.")
}

func printNode(node *store.NodeExtern) {
	// fmt.Println(node.Key[len("/1"):])
	// fmt.Printf("key:%s ttl:%d mod_index:%d\n", node.Key[len("/1"):], node.TTL, node.ModifiedIndex)
}

const (
	revBytesLen            = 8 + 1 + 8
	markedRevBytesLen      = revBytesLen + 1
	markBytePosition       = markedRevBytesLen - 1
	markTombstone     byte = 't'
)

func isTombstone(b []byte) bool {
	return len(b) == markedRevBytesLen && b[markBytePosition] == markTombstone
}

func traverse(st store.Store, dir string) {
	e, err := st.Get(dir, true, false)
	if err != nil {
		panic(err)
	}
	if len(e.Node.Nodes) == 0 {
		st.Delete(dir, true, true)
		return
	}
	for _, node := range e.Node.Nodes {
		if !node.Dir {
			printNode(node)
		} else {
			// fmt.Println(node.Key[len("/1"):])
			traverse(st, node.Key)
		}
	}
}

func rebuild(datadir string) ([]byte, raftpb.HardState, store.Store) {
	waldir := path.Join(datadir, "member", "wal")
	snapdir := path.Join(datadir, "member", "snap")

	ss := snap.New(snapdir)
	snapshot, err := ss.Load()
	if err != nil && err != snap.ErrNoSnapshot {
		panic(err)
	}

	var walsnap walpb.Snapshot
	if snapshot != nil {
		walsnap.Index, walsnap.Term = snapshot.Metadata.Index, snapshot.Metadata.Term
	}

	w, err := wal.OpenForRead(waldir, walsnap)
	if err != nil {
		panic(err)
	}
	defer w.Close()

	meta, hardstate, ents, err := w.ReadAll()
	if err != nil {
		panic(err)
	}

	st := store.New()
	if snapshot != nil {
		err := st.Recovery(snapshot.Data)
		if err != nil {
			panic(err)
		}
	}

	applier := etcdserver.NewApplierV2(st, nil)
	for _, ent := range ents {
		if ent.Type != raftpb.EntryNormal {
			continue
		}

		var raftReq pb.InternalRaftRequest
		if !pbutil.MaybeUnmarshal(&raftReq, ent.Data) { // backward compatible
			var r pb.Request
			pbutil.MustUnmarshal(&r, ent.Data)
			applyRequest(&r, applier)
		} else {
			if raftReq.V2 != nil {
				req := raftReq.V2
				applyRequest(req, applier)
			}
		}
	}

	return meta, hardstate, st
}

func toTTLOptions(r *pb.Request) store.TTLOptionSet {
	refresh, _ := pbutil.GetBool(r.Refresh)
	ttlOptions := store.TTLOptionSet{Refresh: refresh}
	if r.Expiration != 0 {
		ttlOptions.ExpireTime = time.Unix(0, r.Expiration)
	}
	return ttlOptions
}

func applyRequest(r *pb.Request, applyV2 etcdserver.ApplierV2) {
	toTTLOptions(r)
	switch r.Method {
	case "POST":
		applyV2.Post(r)
	case "PUT":
		applyV2.Put(r)
	case "DELETE":
		applyV2.Delete(r)
	case "QGET":
		applyV2.QGet(r)
	case "SYNC":
		applyV2.Sync(r)
	default:
		panic("unknown command")
	}
}
<commit_msg>make 3 members work<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/etcd/etcdserver"
	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
	"github.com/coreos/etcd/mvcc/backend"
	"github.com/coreos/etcd/mvcc/mvccpb"
	"github.com/coreos/etcd/pkg/pbutil"
	"github.com/coreos/etcd/raft/raftpb"
	"github.com/coreos/etcd/snap"
	"github.com/coreos/etcd/store"
	"github.com/coreos/etcd/wal"
	"github.com/coreos/etcd/wal/walpb"
)

var (
	migrateDatadir string
	ttl            time.Duration
)

func init() {
	flag.StringVar(&migrateDatadir, "data-dir", "", "Path to the data directory")
	flag.DurationVar(&ttl, "ttl", time.Hour, "TTL of event keys (default 1 hour)")
	flag.Parse()
}

func main() {
	dbpath := path.Join(migrateDatadir, "member", "snap", "db")

	be := backend.New(dbpath, time.Second, 10000)
	tx := be.BatchTx()

	st := store.New("/0", "/1")
	expireTime := time.Now().Add(ttl)

	tx.Lock()
	err := tx.UnsafeForEach([]byte("key"), func(k, v []byte) error {
		tb := isTombstone(k)

		kv := &mvccpb.KeyValue{}
		kv.Unmarshal(v)

		// This is a compact key..
		if !strings.HasPrefix(string(kv.Key), "/") {
			return nil
		}

		// fmt.Printf("%s %d %d %d\n", string(kv.Key), kv.CreateRevision, kv.ModRevision, kv.Lease)
		ttlOpt := store.TTLOptionSet{}
		if kv.Lease != 0 {
			ttlOpt = store.TTLOptionSet{ExpireTime: expireTime}
		}

		if !tb {
			_, err := st.Set(path.Join("1", string(kv.Key)), false, string(kv.Value), ttlOpt)
			if err != nil {
				return err
			}
		} else {
			st.Delete(string(kv.Key), false, false)
		}

		return nil
	})
	if err != nil {
		panic(err)
	}
	tx.Unlock()

	traverse(st, "/")

	snapshotter := snap.New(path.Join(migrateDatadir, "member", "snap"))
	raftSnap, err := snapshotter.Load()
	if err != nil {
		panic(err)
	}

	metadata, hardstate, oldSt := rebuild(migrateDatadir)

	if err := os.RemoveAll(migrateDatadir); err != nil {
		panic(err)
	}
	if err := os.MkdirAll(path.Join(migrateDatadir, "member", "snap"), 0700); err != nil {
		panic(err)
	}
	walDir := path.Join(migrateDatadir, "member", "wal")

	w, err := wal.Create(walDir, metadata)
	if err != nil {
		panic(err)
	}
	err = w.SaveSnapshot(walpb.Snapshot{Index: hardstate.Commit, Term: hardstate.Term})
	if err != nil {
		panic(err)
	}
	w.Close()

	nodeIDs := []uint64{}
	event, err := oldSt.Get("/0", true, false)
	if err != nil {
		panic(err)
	}
	q := []*store.NodeExtern{}
	q = append(q, event.Node)
	for len(q) > 0 {
		n := q[0]
		q = q[1:]
		v := ""
		if !n.Dir {
			v = *n.Value
		}
		if n.Key != "/0" {
			if n.Key == path.Join("/0", "version") {
				v = "2.3.7"
			}
			if _, err := st.Set(n.Key, n.Dir, v, store.TTLOptionSet{}); err != nil {
				panic(err)
			}

			fields := strings.Split(n.Key, "/")
			if len(fields) == 4 && fields[2] == "members" {
				nodeID, err := strconv.ParseUint(fields[3], 16, 64)
				if err != nil {
					fmt.Println("wrong ID: %s", fields[3])
					panic(err)
				}
				nodeIDs = append(nodeIDs, nodeID)
			}
		}
		for _, next := range n.Nodes {
			q = append(q, next)
		}
	}

	data, err := st.Save()
	if err != nil {
		panic(err)
	}
	raftSnap.Data = data
	raftSnap.Metadata.Index = hardstate.Commit
	raftSnap.Metadata.Term = hardstate.Term
	// fill in nodes by iterating all members in /0 namespace with their IDs
	raftSnap.Metadata.ConfState.Nodes = nodeIDs
	if err := snapshotter.SaveSnap(*raftSnap); err != nil {
		panic(err)
	}
	fmt.Println("Finished.")
}

func printNode(node *store.NodeExtern) {
	// fmt.Println(node.Key[len("/1"):])
	// fmt.Printf("key:%s ttl:%d mod_index:%d\n", node.Key[len("/1"):], node.TTL, node.ModifiedIndex)
}

const (
	revBytesLen            = 8 + 1 + 8
	markedRevBytesLen      = revBytesLen + 1
	markBytePosition       = markedRevBytesLen - 1
	markTombstone     byte = 't'
)

func isTombstone(b []byte) bool {
	return len(b) == markedRevBytesLen && b[markBytePosition] == markTombstone
}

func traverse(st store.Store, dir string) {
	e, err := st.Get(dir, true, false)
	if err != nil {
		panic(err)
	}
	if len(e.Node.Nodes) == 0 {
		st.Delete(dir, true, true)
		return
	}
	for _, node := range e.Node.Nodes {
		if !node.Dir {
			printNode(node)
		} else {
			// fmt.Println(node.Key[len("/1"):])
			traverse(st, node.Key)
		}
	}
}

func rebuild(datadir string) ([]byte, raftpb.HardState, store.Store) {
	waldir := path.Join(datadir, "member", "wal")
	snapdir := path.Join(datadir, "member", "snap")

	ss := snap.New(snapdir)
	snapshot, err := ss.Load()
	if err != nil && err != snap.ErrNoSnapshot {
		panic(err)
	}

	var walsnap walpb.Snapshot
	if snapshot != nil {
		walsnap.Index, walsnap.Term = snapshot.Metadata.Index, snapshot.Metadata.Term
	}

	w, err := wal.OpenForRead(waldir, walsnap)
	if err != nil {
		panic(err)
	}
	defer w.Close()

	meta, hardstate, ents, err := w.ReadAll()
	if err != nil {
		panic(err)
	}

	st := store.New()
	if snapshot != nil {
		err := st.Recovery(snapshot.Data)
		if err != nil {
			panic(err)
		}
	}

	applier := etcdserver.NewApplierV2(st, nil)
	for _, ent := range ents {
		if ent.Type != raftpb.EntryNormal {
			continue
		}

		var raftReq pb.InternalRaftRequest
		if !pbutil.MaybeUnmarshal(&raftReq, ent.Data) { // backward compatible
			var r pb.Request
			pbutil.MustUnmarshal(&r, ent.Data)
			applyRequest(&r, applier)
		} else {
			if raftReq.V2 != nil {
				req := raftReq.V2
				applyRequest(req, applier)
			}
		}
	}

	return meta, hardstate, st
}

func toTTLOptions(r *pb.Request) store.TTLOptionSet {
	refresh, _ := pbutil.GetBool(r.Refresh)
	ttlOptions := store.TTLOptionSet{Refresh: refresh}
	if r.Expiration != 0 {
		ttlOptions.ExpireTime = time.Unix(0, r.Expiration)
	}
	return ttlOptions
}

func applyRequest(r *pb.Request, applyV2 etcdserver.ApplierV2) {
	toTTLOptions(r)
	switch r.Method {
	case "POST":
		applyV2.Post(r)
	case "PUT":
		applyV2.Put(r)
	case "DELETE":
		applyV2.Delete(r)
	case "QGET":
		applyV2.QGet(r)
	case "SYNC":
		applyV2.Sync(r)
	default:
		panic("unknown command")
	}
}
<|endoftext|>


<commit_before>package ratelimit

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"
)

func TestRateLimit(t *testing.T) {
	if os.Getenv("REDIS_URL") == "" {
		t.Skip("skipping redis test since there is no REDIS_URL")
	}

	if time.Now().Minute() > 58 {
		t.Log("Note: The TestRateLimit test is known to have a bug if run near the top of the hour. Since the rate limiter isn't a moving window, it could end up checking against two different buckets on either side of the top of the hour, so if you see that just re-run it after you've passed the top of the hour.")
	}

	rateLimiter := NewRateLimiter(os.Getenv("REDIS_URL"), fmt.Sprintf("worker-test-rl-%d", os.Getpid()))

	ok, err := rateLimiter.RateLimit(context.TODO(), "slow", 2, time.Hour)
	if err != nil {
		t.Fatalf("rate limiter error: %v", err)
	}
	if !ok {
		t.Fatal("expected to not get rate limited, but was limited")
	}

	ok, err = rateLimiter.RateLimit(context.TODO(), "slow", 2, time.Hour)
	if err != nil {
		t.Fatalf("rate limiter error: %v", err)
	}
	if !ok {
		t.Fatal("expected to not get rate limited, but was limited")
	}

	ok, err = rateLimiter.RateLimit("slow", 2, time.Hour)
	if err != nil {
		t.Fatalf("rate limiter error: %v", err)
	}
	if ok {
		t.Fatal("expected to get rate limited, but was not limited")
	}
}
<commit_msg>Pass in last context.TODO to make test pass<commit_after>package ratelimit

import (
	"context"
	"fmt"
	"os"
	"testing"
	"time"
)

func TestRateLimit(t *testing.T) {
	if os.Getenv("REDIS_URL") == "" {
		t.Skip("skipping redis test since there is no REDIS_URL")
	}

	if time.Now().Minute() > 58 {
		t.Log("Note: The TestRateLimit test is known to have a bug if run near the top of the hour. Since the rate limiter isn't a moving window, it could end up checking against two different buckets on either side of the top of the hour, so if you see that just re-run it after you've passed the top of the hour.")
	}

	rateLimiter := NewRateLimiter(os.Getenv("REDIS_URL"), fmt.Sprintf("worker-test-rl-%d", os.Getpid()))

	ok, err := rateLimiter.RateLimit(context.TODO(), "slow", 2, time.Hour)
	if err != nil {
		t.Fatalf("rate limiter error: %v", err)
	}
	if !ok {
		t.Fatal("expected to not get rate limited, but was limited")
	}

	ok, err = rateLimiter.RateLimit(context.TODO(), "slow", 2, time.Hour)
	if err != nil {
		t.Fatalf("rate limiter error: %v", err)
	}
	if !ok {
		t.Fatal("expected to not get rate limited, but was limited")
	}

	ok, err = rateLimiter.RateLimit(context.TODO(), "slow", 2, time.Hour)
	if err != nil {
		t.Fatalf("rate limiter error: %v", err)
	}
	if ok {
		t.Fatal("expected to get rate limited, but was not limited")
	}
}
<|endoftext|>


<commit_before>package alerts

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	util "github.com/MKuranowski/WarsawGTFS/realtime/util"
	gtfsrt "github.com/MobilityData/gtfs-realtime-bindings/golang/gtfs"
	"github.com/golang/protobuf/proto"
)

// Alert contains an internal representation of an alert, which is also marshallable to JSON
type Alert struct {
	ID       string   `json:"id"`
	Routes   []string `json:"routes"`
	Effect   string   `json:"effect"`
	Link     string   `json:"link"`
	Title    string   `json:"title"`
	Body     string   `json:"body"`
	HTMLBody string   `json:"htmlbody"`
}

// alertFromRssItem extracts basic data from an RssItem and puts them into an Alert
func alertFromRssItem(r *rssItem) (a *Alert, err error) {
	a = &Alert{}

	// Extract the ID
	if idMatch := regexID.FindStringIndex(r.GUID); idMatch != nil {
		// Generate a prefix for the ID
		var idPrefix string
		if r.Type == "REDUCED_SERVICE" {
			idPrefix = "A/IMPEDIMENT/"
		} else {
			idPrefix = "A/CHANGE/"
		}

		a.ID = idPrefix + r.GUID[idMatch[0]+3:idMatch[1]]
	} else {
		err = fmt.Errorf("Unable to find alert ID in GUID %q", r.GUID)
		return
	}

	// Extract other data
	a.Effect = r.Type
	a.Link = htmlCleaner.Sanitize(r.Link)
	a.Title = htmlCleaner.Sanitize(r.Description)

	// Extract affected routes from the title
	if strings.Contains(r.Title, ":") {
		routesString := strings.SplitN(r.Title, ":", 2)[1]
		a.Routes = regexRoute.FindAllString(routesString, -1)
	}

	return
}

// makeEntitySelector creates a GTFS-RT []*EntitySelector that "select" applicable routes
func (a *Alert) makeEntitySelector() []*gtfsrt.EntitySelector {
	var entities []*gtfsrt.EntitySelector

	for _, route := range a.Routes {
		entities = append(entities, &gtfsrt.EntitySelector{RouteId: &route})
	}

	return entities
}

// makeAlertEffect creates a GTFS-RT *Alert_Effect that selects the GTFS-RT Effect attribute
func (a *Alert) makeAlertEffect() *gtfsrt.Alert_Effect {
	var effect gtfsrt.Alert_Effect
	if a.Effect == "REDUCED_SERVICE" {
		effect = gtfsrt.Alert_REDUCED_SERVICE
	} else {
		effect = gtfsrt.Alert_OTHER_EFFECT
	}
	return &effect
}

// AsProto returns the Alert marshalled to a GTFS-RT FeedEntity
func (a *Alert) AsProto() *gtfsrt.FeedEntity {
	return &gtfsrt.FeedEntity{
		Id: &a.ID,
		Alert: &gtfsrt.Alert{
			InformedEntity:  a.makeEntitySelector(),
			Effect:          a.makeAlertEffect(),
			Url:             util.MakeTranslatedString(a.Link),
			HeaderText:      util.MakeTranslatedString(a.Title),
			DescriptionText: util.MakeTranslatedString(a.Body),
		},
	}
}

// LoadExternal processes data located on the website saved in a.Link
func (a *Alert) LoadExternal(client exclusiveHTTPClient, routeMap map[string]sort.StringSlice) (err error) {
	doc, err := getWebsite(client, a.Link, a.ID)
	if err != nil {
		return
	}

	// Process flags
	if len(a.Routes) <= 0 {
		flags := getAlertFlags(doc, a.Effect)
		for _, flag := range flags {
			switch flag {
			case "metro":
				a.Routes = append(a.Routes, routeMap["1"]...)
			case "tramwaje":
				a.Routes = append(a.Routes, routeMap["0"]...)
			case "kolej":
			case "skm":
				a.Routes = append(a.Routes, routeMap["2"]...)
			case "autobusy":
				a.Routes = append(a.Routes, routeMap["3"]...)
			}
		}
	}

	// Sanitize the document to make a usable htmlBody
	a.HTMLBody, err = getAlertDesc(doc, a.Effect)
	if err != nil {
		return
	}

	// Strip the cleaned HTMLBody to make a plaintext description
	a.Body, err = getAlertPlaintext(a.HTMLBody)
	return
}

// AlertContainer is a container for multiple alerts, marshallable to JSON
type AlertContainer struct {
	Timestamp time.Time `json:"-"`
	Time      string    `json:"time"`
	Alerts    []*Alert  `json:"alerts"`
}

// AsProto returns this list of alerts marshalled into a GTFS-RT FeedMessage
func (ac *AlertContainer) AsProto() *gtfsrt.FeedMessage {
	msg := util.MakeFeedMessage(ac.Timestamp)
	for _, alert := range ac.Alerts {
		msg.Entity = append(msg.Entity, alert.AsProto())
	}
	return msg
}

// LoadExternal asynchronously calls LoadExternal on all its alerts
func (ac *AlertContainer) LoadExternal(client exclusiveHTTPClient, routeMap map[string]sort.StringSlice, throwErrors bool) error {
	// Make synchronization primitives
	wg := &sync.WaitGroup{}
	errCh := make(chan error, len(ac.Alerts)+1)

	// Make a goroutine for all LoadExternal
	for _, a := range ac.Alerts {
		// Log & update the waitgroup
		log.Printf("Fetching alert desc from %v\n", a.Link)
		wg.Add(1)

		// Call a goroutine to load external data
		go func(a *Alert) {
			defer wg.Done()
			err := a.LoadExternal(client, routeMap)

			// errors are only passed through if requested and if not nil
			if throwErrors && err != nil {
				errCh <- err
			}
		}(a)
	}

	// Wait until all goroutines finished
	wg.Wait()

	// Mark the end of possible errors
	close(errCh)

	// Check if any error was asked to be re-thrown
	for err := range errCh {
		return err
	}
	return nil
}

// Filter removes all alerts without associated routes
func (ac *AlertContainer) Filter() {
	filtered := ac.Alerts[:0]
	// Filter alerts
	for _, a := range ac.Alerts {
		if len(a.Routes) > 0 {
			filtered = append(filtered, a)
		}
	}
	// Garbage collect deleted alerts
	for i := len(filtered); i < len(ac.Alerts); i++ {
		ac.Alerts[i] = nil
	}
	// Set the filtered slice
	ac.Alerts = filtered
}

// SaveJSON marshalls the container into a json file at the given location
func (ac *AlertContainer) SaveJSON(target string) (err error) {
	// Open target file
	f, err := os.Create(target)
	if err != nil {
		return
	}
	defer f.Close()

	// Marshall JSON
	b, err := json.MarshalIndent(ac, "", "  ")
	if err != nil {
		return
	}
	f.Write(b)

	return
}

// SavePB marshalls the container into a GTFS-Realtime protocol buffer file
func (ac *AlertContainer) SavePB(target string, humanReadable bool) (err error) {
	// Open target file
	f, err := os.Create(target)
	if err != nil {
		return
	}
	defer f.Close()

	// Marshall to GTFS-RT
	if humanReadable {
		// Human-readable format
		err = proto.MarshalText(f, ac.AsProto())
		if err != nil {
			return
		}
	} else {
		// Binary format
		var b []byte
		b, err = proto.Marshal(ac.AsProto())
		if err != nil {
			return
		}
		f.Write(b)
	}

	return
}
<commit_msg>alerts: fix invalid entity selectors in gtfs-realtime format<commit_after>package alerts

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	util "github.com/MKuranowski/WarsawGTFS/realtime/util"
	gtfsrt "github.com/MobilityData/gtfs-realtime-bindings/golang/gtfs"
	"github.com/golang/protobuf/proto"
)

// Alert contains an internal representation of an alert, which is also marshallable to JSON
type Alert struct {
	ID       string   `json:"id"`
	Routes   []string `json:"routes"`
	Effect   string   `json:"effect"`
	Link     string   `json:"link"`
	Title    string   `json:"title"`
	Body     string   `json:"body"`
	HTMLBody string   `json:"htmlbody"`
}

// alertFromRssItem extracts basic data from an RssItem and puts them into an Alert
func alertFromRssItem(r *rssItem) (a *Alert, err error) {
	a = &Alert{}

	// Extract the ID
	if idMatch := regexID.FindStringIndex(r.GUID); idMatch != nil {
		// Generate a prefix for the ID
		var idPrefix string
		if r.Type == "REDUCED_SERVICE" {
			idPrefix = "A/IMPEDIMENT/"
		} else {
			idPrefix = "A/CHANGE/"
		}

		a.ID = idPrefix + r.GUID[idMatch[0]+3:idMatch[1]]
	} else {
		err = fmt.Errorf("Unable to find alert ID in GUID %q", r.GUID)
		return
	}

	// Extract other data
	a.Effect = r.Type
	a.Link = htmlCleaner.Sanitize(r.Link)
	a.Title = htmlCleaner.Sanitize(r.Description)

	// Extract affected routes from the title
	if strings.Contains(r.Title, ":") {
		routesString := strings.SplitN(r.Title, ":", 2)[1]
		a.Routes = regexRoute.FindAllString(routesString, -1)
	}

	return
}

// makeEntitySelector creates a GTFS-RT []*EntitySelector that "select" applicable routes
func (a *Alert) makeEntitySelector() []*gtfsrt.EntitySelector {
	var entities []*gtfsrt.EntitySelector

	for i := range a.Routes {
		entities = append(entities, &gtfsrt.EntitySelector{RouteId: &a.Routes[i]})
	}

	return entities
}

// makeAlertEffect creates a GTFS-RT *Alert_Effect that selects the GTFS-RT Effect attribute
func (a *Alert) makeAlertEffect() *gtfsrt.Alert_Effect {
	var effect gtfsrt.Alert_Effect
	if a.Effect == "REDUCED_SERVICE" {
		effect = gtfsrt.Alert_REDUCED_SERVICE
	} else {
		effect = gtfsrt.Alert_OTHER_EFFECT
	}
	return &effect
}

// AsProto returns the Alert marshalled to a GTFS-RT FeedEntity
func (a *Alert) AsProto() *gtfsrt.FeedEntity {
	return &gtfsrt.FeedEntity{
		Id: &a.ID,
		Alert: &gtfsrt.Alert{
			InformedEntity:  a.makeEntitySelector(),
			Effect:          a.makeAlertEffect(),
			Url:             util.MakeTranslatedString(a.Link),
			HeaderText:      util.MakeTranslatedString(a.Title),
			DescriptionText: util.MakeTranslatedString(a.Body),
		},
	}
}

// LoadExternal processes data located on the website saved in a.Link
func (a *Alert) LoadExternal(client exclusiveHTTPClient, routeMap map[string]sort.StringSlice) (err error) {
	doc, err := getWebsite(client, a.Link, a.ID)
	if err != nil {
		return
	}

	// Process flags
	if len(a.Routes) <= 0 {
		flags := getAlertFlags(doc, a.Effect)
		for _, flag := range flags {
			switch flag {
			case "metro":
				a.Routes = append(a.Routes, routeMap["1"]...)
			case "tramwaje":
				a.Routes = append(a.Routes, routeMap["0"]...)
			case "kolej":
			case "skm":
				a.Routes = append(a.Routes, routeMap["2"]...)
			case "autobusy":
				a.Routes = append(a.Routes, routeMap["3"]...)
			}
		}
	}

	// Sanitize the document to make a usable htmlBody
	a.HTMLBody, err = getAlertDesc(doc, a.Effect)
	if err != nil {
		return
	}

	// Strip the cleaned HTMLBody to make a plaintext description
	a.Body, err = getAlertPlaintext(a.HTMLBody)
	return
}

// AlertContainer is a container for multiple alerts, marshallable to JSON
type AlertContainer struct {
	Timestamp time.Time `json:"-"`
	Time      string    `json:"time"`
	Alerts    []*Alert  `json:"alerts"`
}

// AsProto returns this list of alerts marshalled into a GTFS-RT FeedMessage
func (ac *AlertContainer) AsProto() *gtfsrt.FeedMessage {
	msg := util.MakeFeedMessage(ac.Timestamp)
	for _, alert := range ac.Alerts {
		msg.Entity = append(msg.Entity, alert.AsProto())
	}
	return msg
}

// LoadExternal asynchronously calls LoadExternal on all its alerts
func (ac *AlertContainer) LoadExternal(client exclusiveHTTPClient, routeMap map[string]sort.StringSlice, throwErrors bool) error {
	// Make synchronization primitives
	wg := &sync.WaitGroup{}
	errCh := make(chan error, len(ac.Alerts)+1)

	// Make a goroutine for all LoadExternal
	for _, a := range ac.Alerts {
		// Log & update the waitgroup
		log.Printf("Fetching alert desc from %v\n", a.Link)
		wg.Add(1)

		// Call a goroutine to load external data
		go func(a *Alert) {
			defer wg.Done()
			err := a.LoadExternal(client, routeMap)

			// errors are only passed through if requested and if not nil
			if throwErrors && err != nil {
				errCh <- err
			}
		}(a)
	}

	// Wait until all goroutines finished
	wg.Wait()

	// Mark the end of possible errors
	close(errCh)

	// Check if any error was asked to be re-thrown
	for err := range errCh {
		return err
	}
	return nil
}

// Filter removes all alerts without associated routes
func (ac *AlertContainer) Filter() {
	filtered := ac.Alerts[:0]
	// Filter alerts
	for _, a := range ac.Alerts {
		if len(a.Routes) > 0 {
			filtered = append(filtered, a)
		}
	}
	// Garbage collect deleted alerts
	for i := len(filtered); i < len(ac.Alerts); i++ {
		ac.Alerts[i] = nil
	}
	// Set the filtered slice
	ac.Alerts = filtered
}

// SaveJSON marshalls the container into a json file at the given location
func (ac *AlertContainer) SaveJSON(target string) (err error) {
	// Open target file
	f, err := os.Create(target)
	if err != nil {
		return
	}
	defer f.Close()

	// Marshall JSON
	b, err := json.MarshalIndent(ac, "", "  ")
	if err != nil {
		return
	}
	f.Write(b)

	return
}

// SavePB marshalls the container into a GTFS-Realtime protocol buffer file
func (ac *AlertContainer) SavePB(target string, humanReadable bool) (err error) {
	// Open target file
	f, err := os.Create(target)
	if err != nil {
		return
	}
	defer f.Close()

	// Marshall to GTFS-RT
	if humanReadable {
		// Human-readable format
		err = proto.MarshalText(f, ac.AsProto())
		if err != nil {
			return
		}
	} else {
		// Binary format
		var b []byte
		b, err = proto.Marshal(ac.AsProto())
		if err != nil {
			return
		}
		f.Write(b)
	}

	return
}
<|endoftext|>


<commit_before>//  Copyright (c) 2015 Couchbase, Inc.
//  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 rebalance

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
	"testing"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
	"github.com/couchbaselabs/cbgt/rest/monitor"
)

func TestRebalance(t *testing.T) {
	testDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(testDir)

	httpGets := 0
	monitor.HttpGet = func(url string) (resp *http.Response, err error) {
		httpGets++

		return &http.Response{
			StatusCode: 200,
			Body:       ioutil.NopCloser(bytes.NewBuffer([]byte("{}"))),
		}, nil
	}
	defer func() {
		monitor.HttpGet = http.Get
	}()

	tests := []struct {
		label       string
		ops         string // Space separated "+a", "-x".
		params      map[string]string
		expNodes    string // Space separated list of nodes ("a"..."v").
		expIndexes  string // Space separated list of indxes ("x"..."z").
		expStartErr bool
	}{
		{"1st node",
			"+a", nil,
			"a",
			"",
			true,
		},
		{"add 1st index x",
			"+x", nil,
			"a",
			"x",
			false,
		},
		{"add 2nd node b",
			"+b", nil,
			"a b",
			"x",
			false,
		},
		{"remove node b",
			"-b", nil,
			"a",
			"x",
			false,
		},
	}

	cfg := cbgt.NewCfgMem()

	mgrs := map[string]*cbgt.Manager{}

	var mgr0 *cbgt.Manager

	server := "."

	waitUntilEmptyCfgEvents := func(ch chan cbgt.CfgEvent) {
		for {
			select {
			case <-ch:
			default:
				return
			}
		}
	}

	cfgEventsNodeDefsWanted := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.NODE_DEFS_WANTED, cfgEventsNodeDefsWanted)

	waitUntilEmptyCfgEventsNodeDefsWanted := func() {
		waitUntilEmptyCfgEvents(cfgEventsNodeDefsWanted)
	}

	cfgEventsIndexDefs := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.INDEX_DEFS_KEY, cfgEventsIndexDefs)

	waitUntilEmptyCfgEventsIndexDefs := func() {
		waitUntilEmptyCfgEvents(cfgEventsIndexDefs)
	}

	for testi, test := range tests {
		log.Printf("testi: %d, label: %q", testi, test.label)

		for opi, op := range strings.Split(test.ops, " ") {
			log.Printf(" opi: %d, op: %s", opi, op)

			name := op[1:2]

			isIndexOp := name >= "x"
			if isIndexOp {
				indexName := name
				log.Printf(" indexOp: %s, indexName: %s", op[0:1], indexName)

				testCreateIndex(t, mgr0, indexName, test.params,
					waitUntilEmptyCfgEventsIndexDefs)
			} else { // It's a node op.
				nodeName := name
				log.Printf(" nodeOp: %s, nodeName: %s", op[0:1], nodeName)

				register := "wanted"
				if op[0:1] == "-" {
					register = "unknown"
				}
				if test.params["register"] != "" {
					register = test.params["register"]
				}
				if test.params[nodeName+".register"] != "" {
					register = test.params[nodeName+".register"]
				}

				if mgrs[nodeName] != nil {
					mgrs[nodeName].Stop()
					delete(mgrs, nodeName)
				}

				waitUntilEmptyCfgEventsNodeDefsWanted()

				mgr, err := startNodeManager(testDir, cfg,
					name, register, test.params, server)
				if err != nil || mgr == nil {
					t.Errorf("expected no err, got: %#v", err)
				}
				if mgr0 == nil {
					mgr0 = mgr
				}

				if register != "unknown" {
					mgrs[nodeName] = mgr
				}

				mgr.Kick("kick")

				waitUntilEmptyCfgEventsNodeDefsWanted()
			}
		}

		r, err := StartRebalance(cbgt.VERSION, cfg, ".", nil)
		if (test.expStartErr && err == nil) ||
			(!test.expStartErr && err != nil) {
			t.Errorf("testi: %d, label: %q,"+
				" expStartErr: %v, but got: %v",
				testi, test.label,
				test.expStartErr, err)
		}

		if err != nil || r == nil {
			continue
		}

		progressCh := r.ProgressCh()
		if progressCh == nil {
			t.Errorf("expected progressCh")
		}

		err = nil
		for progress := range progressCh {
			if progress.Error != nil {
				err = progress.Error

				log.Printf("saw progress error: %#v\n", progress)
			}
		}

		r.Stop()

		if err != nil {
			t.Errorf("expected no end err, got: %v", err)
		}

		endIndexDefs, endNodeDefs, endPlanPIndexes, endPlanPIndexesCAS, err :=
			cbgt.PlannerGetPlan(cfg, cbgt.VERSION, "")
		if err != nil ||
			endIndexDefs == nil ||
			endNodeDefs == nil ||
			endPlanPIndexes == nil ||
			endPlanPIndexesCAS == 0 {
			t.Errorf("expected no err, got: %#v", err)
		}

		expNodes := strings.Split(test.expNodes, " ")
		if len(expNodes) != len(endNodeDefs.NodeDefs) {
			t.Errorf("len(expNodes) != len(endNodeDefs.NodeDefs), "+
				" expNodes: %#v, endNodeDefs.NodeDefs: %#v",
				expNodes, endNodeDefs.NodeDefs)
		}

		for _, expNode := range expNodes {
			if endNodeDefs.NodeDefs[expNode] == nil {
				t.Errorf("didn't find expNode: %s,"+
					" expNodes: %#v, endNodeDefs.NodeDefs: %#v",
					expNode, expNodes, endNodeDefs.NodeDefs)
			}
		}
	}
}

func testCreateIndex(t *testing.T,
	mgr *cbgt.Manager,
	indexName string,
	params map[string]string,
	waitUntilEmptyCfgEventsIndexDefs func()) {
	sourceType := "primary"
	if params["sourceType"] != "" {
		sourceType = params["sourceType"]
	}
	if params[indexName+".sourceType"] != "" {
		sourceType = params[indexName+".sourceType"]
	}

	sourceName := "default"
	if params["sourceName"] != "" {
		sourceName = params["sourceName"]
	}
	if params[indexName+".sourceName"] != "" {
		sourceName = params[indexName+".sourceName"]
	}

	sourceUUID := ""
	if params["sourceUUID"] != "" {
		sourceUUID = params["sourceUUID"]
	}
	if params[indexName+".sourceUUID"] != "" {
		sourceUUID = params[indexName+".sourceUUID"]
	}

	sourceParams := `{"numPartitions":4}`
	if params["sourceParams"] != "" {
		sourceParams = params["sourceParams"]
	}
	if params[indexName+".sourceParams"] != "" {
		sourceParams = params[indexName+".sourceParams"]
	}

	indexType := "blackhole"
	if params["indexType"] != "" {
		indexType = params["indexType"]
	}
	if params[indexName+".indexType"] != "" {
		indexType = params[indexName+".indexType"]
	}

	indexParams := ""
	if params["indexParams"] != "" {
		indexParams = params["indexParams"]
	}
	if params[indexName+".indexParams"] != "" {
		indexParams = params[indexName+".indexParams"]
	}

	prevIndexUUID := ""
	if params["prevIndexUUID"] != "" {
		prevIndexUUID = params["prevIndexUUID"]
	}
	if params[indexName+".prevIndexUUID"] != "" {
		prevIndexUUID = params[indexName+".prevIndexUUID"]
	}

	planParams := cbgt.PlanParams{
		MaxPartitionsPerPIndex: 1,
	}

	waitUntilEmptyCfgEventsIndexDefs()

	err := mgr.CreateIndex(
		sourceType, sourceName, sourceUUID, sourceParams,
		indexType, indexName, indexParams,
		planParams,
		prevIndexUUID)
	if err != nil {
		t.Errorf("expected no err, got: %#v", err)
	}

	waitUntilEmptyCfgEventsIndexDefs()
}

func startNodeManager(testDir string, cfg cbgt.Cfg, node, register string,
	params map[string]string, server string) (
	mgr *cbgt.Manager, err error) {
	uuid := node
	if params["uuid"] != "" {
		uuid = params["uuid"]
	}
	if params[node+".uuid"] != "" {
		uuid = params[node+".uuid"]
	}

	// No planner in tags because mcp provides the planner.
	tags := []string{"feed", "pindex", "janitor", "queryer"}
	if params["tags"] != "" {
		tags = strings.Split(params["tags"], ",")
	}
	if params[node+".tags"] != "" {
		tags = strings.Split(params[node+".tags"], ",")
	}

	container := ""
	if params["container"] != "" {
		container = params["container"]
	}
	if params[node+".container"] != "" {
		container = params[node+".container"]
	}

	weight := 1
	if params["weight"] != "" {
		weight, err = strconv.Atoi(params["weight"])
	}
	if params[node+".weight"] != "" {
		weight, err = strconv.Atoi(params[node+".weight"])
	}
	if weight < 1 {
		weight = 1
	}

	extras := ""

	bindHttp := node

	dataDir := testDir + string(os.PathSeparator) + node

	os.MkdirAll(dataDir, 0700)

	meh := cbgt.ManagerEventHandlers(nil)

	mgr = cbgt.NewManager(cbgt.VERSION, cfg, uuid,
		tags, container, weight, extras,
		bindHttp, dataDir, server, meh)

	err = mgr.Start(register)
	if err != nil {
		mgr.Stop()

		return nil, err
	}

	return mgr, nil
}
<commit_msg>test that currStates has all the indexes<commit_after>//  Copyright (c) 2015 Couchbase, Inc.
//  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 rebalance

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
	"strings"
	"testing"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
	"github.com/couchbaselabs/cbgt/rest/monitor"
)

func TestRebalance(t *testing.T) {
	testDir, _ := ioutil.TempDir("./tmp", "test")
	defer os.RemoveAll(testDir)

	httpGets := 0
	monitor.HttpGet = func(url string) (resp *http.Response, err error) {
		httpGets++

		return &http.Response{
			StatusCode: 200,
			Body:       ioutil.NopCloser(bytes.NewBuffer([]byte("{}"))),
		}, nil
	}
	defer func() {
		monitor.HttpGet = http.Get
	}()

	tests := []struct {
		label       string
		ops         string // Space separated "+a", "-x".
		params      map[string]string
		expNodes    string // Space separated list of nodes ("a"..."v").
		expIndexes  string // Space separated list of indxes ("x"..."z").
		expStartErr bool
	}{
		{"1st node",
			"+a", nil,
			"a",
			"",
			true,
		},
		{"add 1st index x",
			"+x", nil,
			"a",
			"x",
			false,
		},
		{"add 2nd node b",
			"+b", nil,
			"a b",
			"x",
			false,
		},
		{"remove node b",
			"-b", nil,
			"a",
			"x",
			false,
		},
	}

	cfg := cbgt.NewCfgMem()

	mgrs := map[string]*cbgt.Manager{}

	var mgr0 *cbgt.Manager

	server := "."

	waitUntilEmptyCfgEvents := func(ch chan cbgt.CfgEvent) {
		for {
			select {
			case <-ch:
			default:
				return
			}
		}
	}

	cfgEventsNodeDefsWanted := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.NODE_DEFS_WANTED, cfgEventsNodeDefsWanted)

	waitUntilEmptyCfgEventsNodeDefsWanted := func() {
		waitUntilEmptyCfgEvents(cfgEventsNodeDefsWanted)
	}

	cfgEventsIndexDefs := make(chan cbgt.CfgEvent, 100)
	cfg.Subscribe(cbgt.INDEX_DEFS_KEY, cfgEventsIndexDefs)

	waitUntilEmptyCfgEventsIndexDefs := func() {
		waitUntilEmptyCfgEvents(cfgEventsIndexDefs)
	}

	for testi, test := range tests {
		log.Printf("testi: %d, label: %q", testi, test.label)

		for opi, op := range strings.Split(test.ops, " ") {
			log.Printf(" opi: %d, op: %s", opi, op)

			name := op[1:2]

			isIndexOp := name >= "x"
			if isIndexOp {
				indexName := name
				log.Printf(" indexOp: %s, indexName: %s", op[0:1], indexName)

				testCreateIndex(t, mgr0, indexName, test.params,
					waitUntilEmptyCfgEventsIndexDefs)
			} else { // It's a node op.
				nodeName := name
				log.Printf(" nodeOp: %s, nodeName: %s", op[0:1], nodeName)

				register := "wanted"
				if op[0:1] == "-" {
					register = "unknown"
				}
				if test.params["register"] != "" {
					register = test.params["register"]
				}
				if test.params[nodeName+".register"] != "" {
					register = test.params[nodeName+".register"]
				}

				if mgrs[nodeName] != nil {
					mgrs[nodeName].Stop()
					delete(mgrs, nodeName)
				}

				waitUntilEmptyCfgEventsNodeDefsWanted()

				mgr, err := startNodeManager(testDir, cfg,
					name, register, test.params, server)
				if err != nil || mgr == nil {
					t.Errorf("expected no err, got: %#v", err)
				}
				if mgr0 == nil {
					mgr0 = mgr
				}

				if register != "unknown" {
					mgrs[nodeName] = mgr
				}

				mgr.Kick("kick")

				waitUntilEmptyCfgEventsNodeDefsWanted()
			}
		}

		r, err := StartRebalance(cbgt.VERSION, cfg, ".", nil)
		if (test.expStartErr && err == nil) ||
			(!test.expStartErr && err != nil) {
			t.Errorf("testi: %d, label: %q,"+
				" expStartErr: %v, but got: %v",
				testi, test.label,
				test.expStartErr, err)
		}

		if err != nil || r == nil {
			continue
		}

		progressCh := r.ProgressCh()
		if progressCh == nil {
			t.Errorf("expected progressCh")
		}

		err = nil
		for progress := range progressCh {
			if progress.Error != nil {
				err = progress.Error

				log.Printf("saw progress error: %#v\n", progress)
			}
		}

		r.Stop()

		if err != nil {
			t.Errorf("expected no end err, got: %v", err)
		}

		endIndexDefs, endNodeDefs, endPlanPIndexes, endPlanPIndexesCAS, err :=
			cbgt.PlannerGetPlan(cfg, cbgt.VERSION, "")
		if err != nil ||
			endIndexDefs == nil ||
			endNodeDefs == nil ||
			endPlanPIndexes == nil ||
			endPlanPIndexesCAS == 0 {
			t.Errorf("expected no err, got: %#v", err)
		}

		expNodes := strings.Split(test.expNodes, " ")
		if len(expNodes) != len(endNodeDefs.NodeDefs) {
			t.Errorf("len(expNodes) != len(endNodeDefs.NodeDefs), "+
				" expNodes: %#v, endNodeDefs.NodeDefs: %#v",
				expNodes, endNodeDefs.NodeDefs)
		}

		for _, expNode := range expNodes {
			if endNodeDefs.NodeDefs[expNode] == nil {
				t.Errorf("didn't find expNode: %s,"+
					" expNodes: %#v, endNodeDefs.NodeDefs: %#v",
					expNode, expNodes, endNodeDefs.NodeDefs)
			}
		}

		expIndexes := strings.Split(test.expIndexes, " ")

		r.VisitCurrStates(func(currStates CurrStates) {
			if len(currStates) != len(expIndexes) {
				t.Errorf("len(expIndexes) != len(currStates), "+
					" expIndexes: %#v, currStates: %#v",
					expIndexes, currStates)
			}
		})
	}
}

func testCreateIndex(t *testing.T,
	mgr *cbgt.Manager,
	indexName string,
	params map[string]string,
	waitUntilEmptyCfgEventsIndexDefs func()) {
	sourceType := "primary"
	if params["sourceType"] != "" {
		sourceType = params["sourceType"]
	}
	if params[indexName+".sourceType"] != "" {
		sourceType = params[indexName+".sourceType"]
	}

	sourceName := "default"
	if params["sourceName"] != "" {
		sourceName = params["sourceName"]
	}
	if params[indexName+".sourceName"] != "" {
		sourceName = params[indexName+".sourceName"]
	}

	sourceUUID := ""
	if params["sourceUUID"] != "" {
		sourceUUID = params["sourceUUID"]
	}
	if params[indexName+".sourceUUID"] != "" {
		sourceUUID = params[indexName+".sourceUUID"]
	}

	sourceParams := `{"numPartitions":4}`
	if params["sourceParams"] != "" {
		sourceParams = params["sourceParams"]
	}
	if params[indexName+".sourceParams"] != "" {
		sourceParams = params[indexName+".sourceParams"]
	}

	indexType := "blackhole"
	if params["indexType"] != "" {
		indexType = params["indexType"]
	}
	if params[indexName+".indexType"] != "" {
		indexType = params[indexName+".indexType"]
	}

	indexParams := ""
	if params["indexParams"] != "" {
		indexParams = params["indexParams"]
	}
	if params[indexName+".indexParams"] != "" {
		indexParams = params[indexName+".indexParams"]
	}

	prevIndexUUID := ""
	if params["prevIndexUUID"] != "" {
		prevIndexUUID = params["prevIndexUUID"]
	}
	if params[indexName+".prevIndexUUID"] != "" {
		prevIndexUUID = params[indexName+".prevIndexUUID"]
	}

	planParams := cbgt.PlanParams{
		MaxPartitionsPerPIndex: 1,
	}

	waitUntilEmptyCfgEventsIndexDefs()

	err := mgr.CreateIndex(
		sourceType, sourceName, sourceUUID, sourceParams,
		indexType, indexName, indexParams,
		planParams,
		prevIndexUUID)
	if err != nil {
		t.Errorf("expected no err, got: %#v", err)
	}

	waitUntilEmptyCfgEventsIndexDefs()
}

func startNodeManager(testDir string, cfg cbgt.Cfg, node, register string,
	params map[string]string, server string) (
	mgr *cbgt.Manager, err error) {
	uuid := node
	if params["uuid"] != "" {
		uuid = params["uuid"]
	}
	if params[node+".uuid"] != "" {
		uuid = params[node+".uuid"]
	}

	// No planner in tags because mcp provides the planner.
	tags := []string{"feed", "pindex", "janitor", "queryer"}
	if params["tags"] != "" {
		tags = strings.Split(params["tags"], ",")
	}
	if params[node+".tags"] != "" {
		tags = strings.Split(params[node+".tags"], ",")
	}

	container := ""
	if params["container"] != "" {
		container = params["container"]
	}
	if params[node+".container"] != "" {
		container = params[node+".container"]
	}

	weight := 1
	if params["weight"] != "" {
		weight, err = strconv.Atoi(params["weight"])
	}
	if params[node+".weight"] != "" {
		weight, err = strconv.Atoi(params[node+".weight"])
	}
	if weight < 1 {
		weight = 1
	}

	extras := ""

	bindHttp := node

	dataDir := testDir + string(os.PathSeparator) + node

	os.MkdirAll(dataDir, 0700)

	meh := cbgt.ManagerEventHandlers(nil)

	mgr = cbgt.NewManager(cbgt.VERSION, cfg, uuid,
		tags, container, weight, extras,
		bindHttp, dataDir, server, meh)

	err = mgr.Start(register)
	if err != nil {
		mgr.Stop()

		return nil, err
	}

	return mgr, nil
}
<|endoftext|>


<commit_before>package persona

import (
	"code.google.com/p/gorilla/sessions"
	"net/http"
)

type Store struct {
	store sessions.Store
}

func NewStore(secret string) Store {
	return Store{sessions.NewCookieStore([]byte(secret))}
}

func (s Store) Get(r *http.Request) string {
	session, _ := s.store.Get(r, "session-name")

	if v, ok := session.Values["email"].(string); ok {
		return v
	}

	return ""
}

func (s Store) Set(email string, w http.ResponseWriter, r *http.Request) {
	session, _ := s.store.Get(r, "session-name")
	session.Values["email"] = email
	session.Save(r, w)
}
<commit_msg>Change gorilla/sessions import path to github<commit_after>package persona

import (
	"net/http"

	"github.com/gorilla/sessions"
)

type Store struct {
	store sessions.Store
}

func NewStore(secret string) Store {
	return Store{sessions.NewCookieStore([]byte(secret))}
}

func (s Store) Get(r *http.Request) string {
	session, _ := s.store.Get(r, "session-name")

	if v, ok := session.Values["email"].(string); ok {
		return v
	}

	return ""
}

func (s Store) Set(email string, w http.ResponseWriter, r *http.Request) {
	session, _ := s.store.Get(r, "session-name")
	session.Values["email"] = email
	session.Save(r, w)
}
<|endoftext|>


<commit_before>// Package stow is used to persist objects to a bolt.DB database.
package stow

import (
	"bytes"
	"errors"
	"sync"

	"github.com/boltdb/bolt"
)

var pool *sync.Pool = &sync.Pool{
	New: func() interface{} { return bytes.NewBuffer(nil) },
}

// ErrNotFound indicates object is not in database.
var ErrNotFound = errors.New("not found")

// Store manages objects persistence.
type Store struct {
	db     *bolt.DB
	bucket []byte
	codec  Codec
}

// NewStore creates a new Store, using the underlying
// bolt.DB "bucket" to persist objects.
// NewStore uses GobEncoding, your objects must be registered
// via gob.Register() for this encoding to work.
func NewStore(db *bolt.DB, bucket []byte) *Store {
	return NewCustomStore(db, bucket, GobCodec{})
}

// NewJSONStore creates a new Store, using the underlying
// bolt.DB "bucket" to persist objects as json.
func NewJSONStore(db *bolt.DB, bucket []byte) *Store {
	return NewCustomStore(db, bucket, JSONCodec{})
}

// NewXMLStore creates a new Store, using the underlying
// bolt.DB "bucket" to persist objects as xml.
func NewXMLStore(db *bolt.DB, bucket []byte) *Store {
	return NewCustomStore(db, bucket, XMLCodec{})
}

// NewCustomStore allows you to create a store with
// a custom underlying Encoding
func NewCustomStore(db *bolt.DB, bucket []byte, codec Codec) *Store {
	return &Store{db: db, bucket: bucket, codec: codec}
}

func (s *Store) marshal(val interface{}) (data []byte, err error) {
	buf := pool.Get().(*bytes.Buffer)
	err = s.codec.NewEncoder(buf).Encode(val)
	data = append(data, buf.Bytes()...)
	buf.Reset()
	pool.Put(buf)

	return data, err
}

func (s *Store) unmarshal(data []byte, val interface{}) (err error) {
	return s.codec.NewDecoder(bytes.NewReader(data)).Decode(val)
}

func (s *Store) toBytes(key interface{}) (keyBytes []byte, err error) {
	switch k := key.(type) {
	case string:
		return []byte(k), nil
	case []byte:
		return k, nil
	default:
		return s.marshal(key)
	}
}

// Put will store b with key "key". If key is []byte or string it uses the key
// directly. Otherwise, it marshals the given type into bytes using the stores Encoder.
func (s *Store) Put(key interface{}, b interface{}) error {
	keyBytes, err := s.toBytes(key)
	if err != nil {
		return err
	}
	return s.put(keyBytes, b)
}

// Put will store b with key "key". If key is []byte or string it uses the key
// directly. Otherwise, it marshals the given type into bytes using the stores Encoder.
func (s *Store) put(key []byte, b interface{}) (err error) {
	var data []byte
	data, err = s.marshal(b)

	return s.db.Update(func(tx *bolt.Tx) error {
		objects, err := tx.CreateBucketIfNotExists(s.bucket)
		if err != nil {
			return err
		}
		objects.Put(key, data)
		return nil
	})
}

// Pull will retrieve b with key "key", and removes it from the store.
func (s *Store) Pull(key interface{}, b interface{}) error {
	keyBytes, err := s.toBytes(key)
	if err != nil {
		return err
	}
	return s.pull(keyBytes, b)
}

// Pull will retrieve b with key "key", and removes it from the store.
func (s *Store) pull(key []byte, b interface{}) error {
	buf := pool.Get().(*bytes.Buffer)
	defer func() {
		buf.Reset()
		pool.Put(buf)
	}()

	err := s.db.Update(func(tx *bolt.Tx) error {
		objects := tx.Bucket(s.bucket)
		if objects == nil {
			return ErrNotFound
		}

		data := objects.Get(key)
		if data == nil {
			return ErrNotFound
		}

		buf.Write(data)
		objects.Delete(key)
		return nil
	})

	if err != nil {
		return err
	}

	return s.unmarshal(buf.Bytes(), b)
}

// Get will retreive b with key "key". If key is []byte or string it uses the key
// directly. Otherwise, it marshals the given type into bytes using the stores Encoder.
func (s *Store) Get(key interface{}, b interface{}) error {
	keyBytes, err := s.toBytes(key)
	if err != nil {
		return err
	}
	return s.get(keyBytes, b)
}

// Get will retreive b with key "key"
func (s *Store) get(key []byte, b interface{}) error {
	buf := bytes.NewBuffer(nil)
	err := s.db.Update(func(tx *bolt.Tx) error {
		objects := tx.Bucket(s.bucket)
		if objects == nil {
			return ErrNotFound
		}
		data := objects.Get(key)
		if data == nil {
			return ErrNotFound
		}
		buf.Write(data)
		return nil
	})

	if err != nil {
		return err
	}

	return s.unmarshal(buf.Bytes(), b)
}

// ForEach will run do on each object in the store.
// do can be a function which takes either: 1 param which will take on each "value"
// or 2 params where the first param is the "key" and the second is the "value".
func (s *Store) ForEach(do interface{}) error {
	fc, err := newFuncCall(s, do)
	if err != nil {
		return err
	}

	return s.db.Update(func(tx *bolt.Tx) error {
		objects := tx.Bucket(s.bucket)
		if objects == nil {
			return nil
		}
		return objects.ForEach(fc.call)
	})
}

// DeleteAll empties the store
func (s *Store) DeleteAll() error {
	return s.db.Update(func(tx *bolt.Tx) error {
		return tx.DeleteBucket(s.bucket)
	})
}
<commit_msg>Use View instead of Update with ForEach and Get<commit_after>// Package stow is used to persist objects to a bolt.DB database.
package stow

import (
	"bytes"
	"errors"
	"sync"

	"github.com/boltdb/bolt"
)

var pool *sync.Pool = &sync.Pool{
	New: func() interface{} { return bytes.NewBuffer(nil) },
}

// ErrNotFound indicates object is not in database.
var ErrNotFound = errors.New("not found")

// Store manages objects persistence.
type Store struct {
	db     *bolt.DB
	bucket []byte
	codec  Codec
}

// NewStore creates a new Store, using the underlying
// bolt.DB "bucket" to persist objects.
// NewStore uses GobEncoding, your objects must be registered
// via gob.Register() for this encoding to work.
func NewStore(db *bolt.DB, bucket []byte) *Store {
	return NewCustomStore(db, bucket, GobCodec{})
}

// NewJSONStore creates a new Store, using the underlying
// bolt.DB "bucket" to persist objects as json.
func NewJSONStore(db *bolt.DB, bucket []byte) *Store {
	return NewCustomStore(db, bucket, JSONCodec{})
}

// NewXMLStore creates a new Store, using the underlying
// bolt.DB "bucket" to persist objects as xml.
func NewXMLStore(db *bolt.DB, bucket []byte) *Store {
	return NewCustomStore(db, bucket, XMLCodec{})
}

// NewCustomStore allows you to create a store with
// a custom underlying Encoding
func NewCustomStore(db *bolt.DB, bucket []byte, codec Codec) *Store {
	return &Store{db: db, bucket: bucket, codec: codec}
}

func (s *Store) marshal(val interface{}) (data []byte, err error) {
	buf := pool.Get().(*bytes.Buffer)
	err = s.codec.NewEncoder(buf).Encode(val)
	data = append(data, buf.Bytes()...)
	buf.Reset()
	pool.Put(buf)

	return data, err
}

func (s *Store) unmarshal(data []byte, val interface{}) (err error) {
	return s.codec.NewDecoder(bytes.NewReader(data)).Decode(val)
}

func (s *Store) toBytes(key interface{}) (keyBytes []byte, err error) {
	switch k := key.(type) {
	case string:
		return []byte(k), nil
	case []byte:
		return k, nil
	default:
		return s.marshal(key)
	}
}

// Put will store b with key "key". If key is []byte or string it uses the key
// directly. Otherwise, it marshals the given type into bytes using the stores Encoder.
func (s *Store) Put(key interface{}, b interface{}) error {
	keyBytes, err := s.toBytes(key)
	if err != nil {
		return err
	}
	return s.put(keyBytes, b)
}

// Put will store b with key "key". If key is []byte or string it uses the key
// directly. Otherwise, it marshals the given type into bytes using the stores Encoder.
func (s *Store) put(key []byte, b interface{}) (err error) {
	var data []byte
	data, err = s.marshal(b)

	return s.db.Update(func(tx *bolt.Tx) error {
		objects, err := tx.CreateBucketIfNotExists(s.bucket)
		if err != nil {
			return err
		}
		objects.Put(key, data)
		return nil
	})
}

// Pull will retrieve b with key "key", and removes it from the store.
func (s *Store) Pull(key interface{}, b interface{}) error {
	keyBytes, err := s.toBytes(key)
	if err != nil {
		return err
	}
	return s.pull(keyBytes, b)
}

// Pull will retrieve b with key "key", and removes it from the store.
func (s *Store) pull(key []byte, b interface{}) error {
	buf := pool.Get().(*bytes.Buffer)
	defer func() {
		buf.Reset()
		pool.Put(buf)
	}()

	err := s.db.Update(func(tx *bolt.Tx) error {
		objects := tx.Bucket(s.bucket)
		if objects == nil {
			return ErrNotFound
		}

		data := objects.Get(key)
		if data == nil {
			return ErrNotFound
		}

		buf.Write(data)
		objects.Delete(key)
		return nil
	})

	if err != nil {
		return err
	}

	return s.unmarshal(buf.Bytes(), b)
}

// Get will retreive b with key "key". If key is []byte or string it uses the key
// directly. Otherwise, it marshals the given type into bytes using the stores Encoder.
func (s *Store) Get(key interface{}, b interface{}) error {
	keyBytes, err := s.toBytes(key)
	if err != nil {
		return err
	}
	return s.get(keyBytes, b)
}

// Get will retreive b with key "key"
func (s *Store) get(key []byte, b interface{}) error {
	buf := bytes.NewBuffer(nil)
	err := s.db.View(func(tx *bolt.Tx) error {
		objects := tx.Bucket(s.bucket)
		if objects == nil {
			return ErrNotFound
		}
		data := objects.Get(key)
		if data == nil {
			return ErrNotFound
		}
		buf.Write(data)
		return nil
	})

	if err != nil {
		return err
	}

	return s.unmarshal(buf.Bytes(), b)
}

// ForEach will run do on each object in the store.
// do can be a function which takes either: 1 param which will take on each "value"
// or 2 params where the first param is the "key" and the second is the "value".
func (s *Store) ForEach(do interface{}) error {
	fc, err := newFuncCall(s, do)
	if err != nil {
		return err
	}

	return s.db.View(func(tx *bolt.Tx) error {
		objects := tx.Bucket(s.bucket)
		if objects == nil {
			return nil
		}
		return objects.ForEach(fc.call)
	})
}

// DeleteAll empties the store
func (s *Store) DeleteAll() error {
	return s.db.Update(func(tx *bolt.Tx) error {
		return tx.DeleteBucket(s.bucket)
	})
}
<|endoftext|>


<commit_before>package ecql

import (
	"strings"

	"fmt"
)

type queryType int

const (
	selectQuery = iota
	insertQuery
	deleteQuery
	updateQuery
	countQuery
)

// Table contains the information of a table in cassandra.
type Table struct {
	Name       string
	KeyColumns []string
	Columns    []Column
}

// Column contains the information of a column in a table required
// to create a map for it.  Every element of position represents an
// anonymous nesting--with a single position representing an immediate named field
type Column struct {
	Name     string
	Position []int
}

func (t *Table) BuildQuery(qt queryType) (string, error) {
	var cql string
	switch qt {
	case selectQuery:
		cql = fmt.Sprintf("SELECT %s FROM %s WHERE %s", t.getCols(), t.Name, appendCols(t.KeyColumns))
	case insertQuery:
		cql = fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s)", t.Name, t.getCols(), t.getQms())
	case deleteQuery:
		cql = fmt.Sprintf("DELETE FROM %s WHERE %s", t.Name, appendCols(t.KeyColumns))
	case updateQuery:
		// cql = "UPDATE %s WHERE %s = ?"
		return "", ErrInvalidQueryType
	case countQuery:
		cql = fmt.Sprintf("SELECT COUNT(1) FROM %s WHERE %s", t.Name, appendCols(t.KeyColumns))
	default:
		return "", ErrInvalidQueryType
	}

	return cql, nil
}

func (t *Table) getCols() string {
	names := make([]string, len(t.Columns))
	for i := range t.Columns {
		names[i] = t.Columns[i].Name
	}
	return strings.Join(names, ",")
}

func (t *Table) getQms() string {
	return qms(len(t.Columns))
}

func appendCols(cols []string) string {
	parts := make([]string, len(cols))
	for i := range cols {
		parts[i] = fmt.Sprintf("%s = ?", cols[i])
	}
	return strings.Join(parts, " AND ")
}
<commit_msg>matt --better explanation for position<commit_after>package ecql

import (
	"strings"

	"fmt"
)

type queryType int

const (
	selectQuery = iota
	insertQuery
	deleteQuery
	updateQuery
	countQuery
)

// Table contains the information of a table in cassandra.
type Table struct {
	Name       string
	KeyColumns []string
	Columns    []Column
}

// Column contains the information of a column in a table required
// to create a map for it.
// Every element of position represents its order in a hierarchy of nested structs
type Column struct {
	Name     string
	Position []int
}

func (t *Table) BuildQuery(qt queryType) (string, error) {
	var cql string
	switch qt {
	case selectQuery:
		cql = fmt.Sprintf("SELECT %s FROM %s WHERE %s", t.getCols(), t.Name, appendCols(t.KeyColumns))
	case insertQuery:
		cql = fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s)", t.Name, t.getCols(), t.getQms())
	case deleteQuery:
		cql = fmt.Sprintf("DELETE FROM %s WHERE %s", t.Name, appendCols(t.KeyColumns))
	case updateQuery:
		// cql = "UPDATE %s WHERE %s = ?"
		return "", ErrInvalidQueryType
	case countQuery:
		cql = fmt.Sprintf("SELECT COUNT(1) FROM %s WHERE %s", t.Name, appendCols(t.KeyColumns))
	default:
		return "", ErrInvalidQueryType
	}

	return cql, nil
}

func (t *Table) getCols() string {
	names := make([]string, len(t.Columns))
	for i := range t.Columns {
		names[i] = t.Columns[i].Name
	}
	return strings.Join(names, ",")
}

func (t *Table) getQms() string {
	return qms(len(t.Columns))
}

func appendCols(cols []string) string {
	parts := make([]string, len(cols))
	for i := range cols {
		parts[i] = fmt.Sprintf("%s = ?", cols[i])
	}
	return strings.Join(parts, " AND ")
}
<|endoftext|>


<commit_before>// Copyright 2017 Zack Guo <zack.y.guo@gmail.com>. All rights reserved.
// Use of this source code is governed by a MIT license that can
// be found in the LICENSE file.

package termui

import "strings"

/* Table is like:

┌Awesome Table ────────────────────────────────────────────────┐
│  Col0          | Col1 | Col2 | Col3  | Col4  | Col5  | Col6  |
│──────────────────────────────────────────────────────────────│
│  Some Item #1  | AAA  | 123  | CCCCC | EEEEE | GGGGG | IIIII |
│──────────────────────────────────────────────────────────────│
│  Some Item #2  | BBB  | 456  | DDDDD | FFFFF | HHHHH | JJJJJ |
└──────────────────────────────────────────────────────────────┘

Datapoints are a two dimensional array of strings: [][]string

Example:
	data := [][]string{
		{"Col0", "Col1", "Col3", "Col4", "Col5", "Col6"},
		{"Some Item #1", "AAA", "123", "CCCCC", "EEEEE", "GGGGG", "IIIII"},
		{"Some Item #2", "BBB", "456", "DDDDD", "FFFFF", "HHHHH", "JJJJJ"},
	}

	table := termui.NewTable()
	table.Rows = data  // type [][]string
	table.FgColor = termui.ColorWhite
	table.BgColor = termui.ColorDefault
	table.Height = 7
	table.Width = 62
	table.Y = 0
	table.X = 0
	table.Border = true
*/

// Table tracks all the attributes of a Table instance
type Table struct {
	Block
	Rows      [][]string
	CellWidth []int
	FgColor   Attribute
	BgColor   Attribute
	FgColors  []Attribute
	BgColors  []Attribute
	Separator bool
	TextAlign Align
}

// NewTable returns a new Table instance
func NewTable() *Table {
	table := &Table{Block: *NewBlock()}
	table.FgColor = ColorWhite
	table.BgColor = ColorDefault
	table.Separator = true
	return table
}

// CellsWidth calculates the width of a cell array and returns an int
func cellsWidth(cells []Cell) int {
	width := 0
	for _, c := range cells {
		width += c.Width()
	}
	return width
}

// Analysis generates and returns an array of []Cell that represent all columns in the Table
func (table *Table) Analysis() [][]Cell {
	var rowCells [][]Cell
	length := len(table.Rows)
	if length < 1 {
		return rowCells
	}

	if len(table.FgColors) == 0 {
		table.FgColors = make([]Attribute, len(table.Rows))
	}
	if len(table.BgColors) == 0 {
		table.BgColors = make([]Attribute, len(table.Rows))
	}

	cellWidths := make([]int, len(table.Rows[0]))

	for y, row := range table.Rows {
		if table.FgColors[y] == 0 {
			table.FgColors[y] = table.FgColor
		}
		if table.BgColors[y] == 0 {
			table.BgColors[y] = table.BgColor
		}
		for x, str := range row {
			cells := DefaultTxBuilder.Build(str, table.FgColors[y], table.BgColors[y])
			cw := cellsWidth(cells)
			if cellWidths[x] < cw {
				cellWidths[x] = cw
			}
			rowCells = append(rowCells, cells)
		}
	}
	table.CellWidth = cellWidths
	return rowCells
}

// SetSize calculates the table size and sets the internal value
func (table *Table) SetSize() {
	length := len(table.Rows)
	if table.Separator {
		table.Height = length*2 + 1
	} else {
		table.Height = length + 2
	}
	table.Width = 2
	if length != 0 {
		for _, cellWidth := range table.CellWidth {
			table.Width += cellWidth + 3
		}
	}
}

// CalculatePosition ...
func (table *Table) CalculatePosition(x int, y int, coordinateX *int, coordinateY *int, cellStart *int) {
	if table.Separator {
		*coordinateY = table.innerArea.Min.Y + y*2
	} else {
		*coordinateY = table.innerArea.Min.Y + y
	}
	if x == 0 {
		*cellStart = table.innerArea.Min.X
	} else {
		*cellStart += table.CellWidth[x-1] + 3
	}

	switch table.TextAlign {
	case AlignRight:
		*coordinateX = *cellStart + (table.CellWidth[x] - len(table.Rows[y][x])) + 2
	case AlignCenter:
		*coordinateX = *cellStart + (table.CellWidth[x]-len(table.Rows[y][x]))/2 + 2
	default:
		*coordinateX = *cellStart + 2
	}
}

// Buffer ...
func (table *Table) Buffer() Buffer {
	buffer := table.Block.Buffer()
	rowCells := table.Analysis()
	pointerX := table.innerArea.Min.X + 2
	pointerY := table.innerArea.Min.Y
	borderPointerX := table.innerArea.Min.X
	for y, row := range table.Rows {
		for x := range row {
			table.CalculatePosition(x, y, &pointerX, &pointerY, &borderPointerX)
			background := DefaultTxBuilder.Build(strings.Repeat(" ", table.CellWidth[x]+3), table.BgColors[y], table.BgColors[y])
			cells := rowCells[y*len(row)+x]
			for i, back := range background {
				buffer.Set(borderPointerX+i, pointerY, back)
			}

			coordinateX := pointerX
			for _, printer := range cells {
				buffer.Set(coordinateX, pointerY, printer)
				coordinateX += printer.Width()
			}

			if x != 0 {
				dividors := DefaultTxBuilder.Build("|", table.FgColors[y], table.BgColors[y])
				for _, dividor := range dividors {
					buffer.Set(borderPointerX, pointerY, dividor)
				}
			}
		}

		if table.Separator {
			border := DefaultTxBuilder.Build(strings.Repeat("─", table.Width-2), table.FgColor, table.BgColor)
			for i, cell := range border {
				buffer.Set(i+1, pointerY+1, cell)
			}
		}
	}

	return buffer
}
<commit_msg>Fix table horizontal separator bug.<commit_after>// Copyright 2017 Zack Guo <zack.y.guo@gmail.com>. All rights reserved.
// Use of this source code is governed by a MIT license that can
// be found in the LICENSE file.

package termui

import "strings"

/* Table is like:

┌Awesome Table ────────────────────────────────────────────────┐
│  Col0          | Col1 | Col2 | Col3  | Col4  | Col5  | Col6  |
│──────────────────────────────────────────────────────────────│
│  Some Item #1  | AAA  | 123  | CCCCC | EEEEE | GGGGG | IIIII |
│──────────────────────────────────────────────────────────────│
│  Some Item #2  | BBB  | 456  | DDDDD | FFFFF | HHHHH | JJJJJ |
└──────────────────────────────────────────────────────────────┘

Datapoints are a two dimensional array of strings: [][]string

Example:
	data := [][]string{
		{"Col0", "Col1", "Col3", "Col4", "Col5", "Col6"},
		{"Some Item #1", "AAA", "123", "CCCCC", "EEEEE", "GGGGG", "IIIII"},
		{"Some Item #2", "BBB", "456", "DDDDD", "FFFFF", "HHHHH", "JJJJJ"},
	}

	table := termui.NewTable()
	table.Rows = data  // type [][]string
	table.FgColor = termui.ColorWhite
	table.BgColor = termui.ColorDefault
	table.Height = 7
	table.Width = 62
	table.Y = 0
	table.X = 0
	table.Border = true
*/

// Table tracks all the attributes of a Table instance
type Table struct {
	Block
	Rows      [][]string
	CellWidth []int
	FgColor   Attribute
	BgColor   Attribute
	FgColors  []Attribute
	BgColors  []Attribute
	Separator bool
	TextAlign Align
}

// NewTable returns a new Table instance
func NewTable() *Table {
	table := &Table{Block: *NewBlock()}
	table.FgColor = ColorWhite
	table.BgColor = ColorDefault
	table.Separator = true
	return table
}

// CellsWidth calculates the width of a cell array and returns an int
func cellsWidth(cells []Cell) int {
	width := 0
	for _, c := range cells {
		width += c.Width()
	}
	return width
}

// Analysis generates and returns an array of []Cell that represent all columns in the Table
func (table *Table) Analysis() [][]Cell {
	var rowCells [][]Cell
	length := len(table.Rows)
	if length < 1 {
		return rowCells
	}

	if len(table.FgColors) == 0 {
		table.FgColors = make([]Attribute, len(table.Rows))
	}
	if len(table.BgColors) == 0 {
		table.BgColors = make([]Attribute, len(table.Rows))
	}

	cellWidths := make([]int, len(table.Rows[0]))

	for y, row := range table.Rows {
		if table.FgColors[y] == 0 {
			table.FgColors[y] = table.FgColor
		}
		if table.BgColors[y] == 0 {
			table.BgColors[y] = table.BgColor
		}
		for x, str := range row {
			cells := DefaultTxBuilder.Build(str, table.FgColors[y], table.BgColors[y])
			cw := cellsWidth(cells)
			if cellWidths[x] < cw {
				cellWidths[x] = cw
			}
			rowCells = append(rowCells, cells)
		}
	}
	table.CellWidth = cellWidths
	return rowCells
}

// SetSize calculates the table size and sets the internal value
func (table *Table) SetSize() {
	length := len(table.Rows)
	if table.Separator {
		table.Height = length*2 + 1
	} else {
		table.Height = length + 2
	}
	table.Width = 2
	if length != 0 {
		for _, cellWidth := range table.CellWidth {
			table.Width += cellWidth + 3
		}
	}
}

// CalculatePosition ...
func (table *Table) CalculatePosition(x int, y int, coordinateX *int, coordinateY *int, cellStart *int) {
	if table.Separator {
		*coordinateY = table.innerArea.Min.Y + y*2
	} else {
		*coordinateY = table.innerArea.Min.Y + y
	}
	if x == 0 {
		*cellStart = table.innerArea.Min.X
	} else {
		*cellStart += table.CellWidth[x-1] + 3
	}

	switch table.TextAlign {
	case AlignRight:
		*coordinateX = *cellStart + (table.CellWidth[x] - len(table.Rows[y][x])) + 2
	case AlignCenter:
		*coordinateX = *cellStart + (table.CellWidth[x]-len(table.Rows[y][x]))/2 + 2
	default:
		*coordinateX = *cellStart + 2
	}
}

// Buffer ...
func (table *Table) Buffer() Buffer {
	buffer := table.Block.Buffer()
	rowCells := table.Analysis()
	pointerX := table.innerArea.Min.X + 2
	pointerY := table.innerArea.Min.Y
	borderPointerX := table.innerArea.Min.X
	for y, row := range table.Rows {
		for x := range row {
			table.CalculatePosition(x, y, &pointerX, &pointerY, &borderPointerX)
			background := DefaultTxBuilder.Build(strings.Repeat(" ", table.CellWidth[x]+3), table.BgColors[y], table.BgColors[y])
			cells := rowCells[y*len(row)+x]
			for i, back := range background {
				buffer.Set(borderPointerX+i, pointerY, back)
			}

			coordinateX := pointerX
			for _, printer := range cells {
				buffer.Set(coordinateX, pointerY, printer)
				coordinateX += printer.Width()
			}

			if x != 0 {
				dividors := DefaultTxBuilder.Build("|", table.FgColors[y], table.BgColors[y])
				for _, dividor := range dividors {
					buffer.Set(borderPointerX, pointerY, dividor)
				}
			}
		}

		if table.Separator {
			border := DefaultTxBuilder.Build(strings.Repeat("─", table.Width-2), table.FgColor, table.BgColor)
			for i, cell := range border {
				buffer.Set(table.innerArea.Min.X+i, pointerY+1, cell)
			}
		}
	}

	return buffer
}
<|endoftext|>


<commit_before>// pushk pushes a new version of an app.
//
// pushk
// pushk docserver
// pushk --rollback docserver
// pushk --project=skia-public docserver
// pushk --rollback --project=skia-public docserver
// pushk --dry-run my-new-service
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/gcr"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	repoUrlTemplate = "https://skia.googlesource.com/%s-config"
	repoBaseDir     = "/tmp"
	repoDirTemplate = "/tmp/%s-config"
)

func init() {
	flag.Usage = func() {
		fmt.Printf("Usage: pushk <flags> [zero or more image names]\n\n")
		fmt.Printf(`pushk pushes a new version of an app.

The command:
  1. Modifies the kubernetes yaml files with the new image.
  2. Commits the changes to the config repo.
  3. Applies the changes with kubectl.

The config is stored in a separate repo that will automaticaly be checked out
under /tmp.

The command applies the changes by default, or just changes the local yaml files
if --dry-run is supplied.

If no image names are supplied then pushk looks through all the yaml files for
appropriate images (ones that match the SERVER and project) and tries to push a
new image for each of them.

Examples:
  # Push the latest version of all images from the given container repository.
  pushk

  # Pusk an exact tag.
  pushk gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07

  # Push the latest version of docserver.
  pushk docserver --message="Fix bug #1234"

  # Push the latest version of docserver and iap-proxy
  pushk docserver iap-proxy

  # Rollback docserver.
  pushk --rollback docserver

  # Compute any changes a push to docserver will make, but do not apply them.
  pushk --dry-run docserver

`)
		flag.PrintDefaults()
	}
}

// toFullRepoURL converts the project name into a git repo URL.
func toFullRepoURL(s string) string {
	return fmt.Sprintf(repoUrlTemplate, s)

}

// toRepoDir converts the project name into a git repo directory name.
func toRepoDir(s string) string {
	return fmt.Sprintf(repoDirTemplate, s)
}

// flags
var (
	dryRun   = flag.Bool("dry-run", false, "If true then do not run the kubectl command to apply the changes, and do not commit the changes to the config repo.")
	message  = flag.String("message", "Push", "Message to go along with the change.")
	project  = flag.String("project", "skia-public", "The GCE project name.")
	rollback = flag.Bool("rollback", false, "If true go back to the second most recent image, otherwise use most recent image.")
)

var (
	validTag = regexp.MustCompile(`^\d\d\d\d-\d\d-\d\dT\d\d_\d\d_\d\dZ-.+$`)
)

// filter strips the list of tags down to only the ones that conform to our
// constraints and also checks that there are enough tags. The results
// are sorted in ascending order, so oldest tags are first, newest tags
// are last.
func filter(tags []string) ([]string, error) {
	validTags := []string{}
	for _, t := range tags {
		if validTag.MatchString(t) {
			validTags = append(validTags, t)
		}
	}
	sort.Strings(validTags)
	if len(validTags) == 0 {
		return nil, fmt.Errorf("Not enough tags returned.")
	}
	return validTags, nil
}

// findAllImageNames searches for all the images that come from the given
// project container registry across all the yaml files listed in filenames.
func findAllImageNames(filenames []string, server, project string) []string {
	// allImageRegex has the following groups returned on match:
	// 0 - the entire line
	// 1 - the image name
	allImageRegex := regexp.MustCompile(fmt.Sprintf(`(?m)^\s+image:\s+%s/%s/([^:]+):\S+\s*$`, server, project))
	filenameSet := util.StringSet{}
	for _, filename := range filenames {
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			sklog.Errorf("Failed to read %q (skipping): %s", filename, err)
			continue
		}
		matches := allImageRegex.FindAllStringSubmatch(string(b), -1)
		for _, m := range matches {
			filenameSet[m[1]] = true
		}
	}
	return filenameSet.Keys()
}

// tagProvider is a type that returns the correct tag to push for the given imageName.
type tagProvider func(imageName string) ([]string, error)

// imageFromCmdLineImage handles image names, which can be either short, ala 'fiddler', or exact,
// such as gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07, both of which can
// appear on the command-line.
func imageFromCmdLineImage(imageName string, tp tagProvider) (string, error) {
	if strings.HasPrefix(imageName, "gcr.io/") {
		if *rollback {
			return "", fmt.Errorf("Supplying a fully qualified image name and the --rollback flag are mutually exclusive.")
		}
		return imageName, nil
	}
	// Get all the tags for the selected image.
	tags, err := tp(imageName)
	if err != nil {
		return "", fmt.Errorf("Tag provider failed: %s", err)
	}

	// Filter the tags
	tags, err = filter(tags)
	if err != nil {
		return "", fmt.Errorf("Failed to filter: %s", err)
	}

	// Pick the target tag we want to move to.
	tag := tags[len(tags)-1]
	if *rollback {
		if len(tags) < 2 {
			return "", fmt.Errorf("No version to rollback to.")
		}
		tag = tags[len(tags)-2]
	}

	// The full docker image name and tag of the image we want to deploy.
	return fmt.Sprintf("%s/%s/%s:%s", gcr.SERVER, *project, imageName, tag), nil
}

func main() {
	common.Init()

	ctx := context.Background()
	repoDir := toRepoDir(*project)
	checkout, err := git.NewCheckout(ctx, toFullRepoURL(*project), repoBaseDir)
	if err != nil {
		sklog.Fatalf("Failed to check out config repo: %s", err)
	}
	if err := checkout.Update(ctx); err != nil {
		sklog.Fatalf("Failed to update repo: %s", err)
	}

	// Get all the yaml files.
	filenames, err := filepath.Glob(filepath.Join(repoDir, "*.yaml"))
	if err != nil {
		sklog.Fatal(err)
	}

	tokenSource := auth.NewGCloudTokenSource(*project)
	imageNames := flag.Args()
	if len(imageNames) == 0 {
		imageNames = findAllImageNames(filenames, gcr.SERVER, *project)
		if len(imageNames) == 0 {
			fmt.Printf("Failed to find any images that match kubernetes directory: %q and project: %q.", repoDir, *project)
			flag.Usage()
			os.Exit(1)
		}
	}
	sklog.Infof("Pushing the following images: %q", imageNames)

	gcrTagProvider := func(imageName string) ([]string, error) {
		return gcr.NewClient(tokenSource, *project, imageName).Tags()
	}

	changed := util.StringSet{}
	for _, imageName := range imageNames {
		image, err := imageFromCmdLineImage(imageName, gcrTagProvider)
		if err != nil {
			sklog.Fatal(err)
		}

		// imageRegex has the following groups returned on match:
		// 0 - the entire line
		// 1 - the prefix, i.e. image:, with correct spacing.
		// 2 - full image name
		// 3 - just the tag
		//
		// We pull out the 'prefix' so we can use it when
		// we rewrite the image: line so the indent level is
		// unchanged.
		parts := strings.SplitN(image, ":", 2)
		if len(parts) != 2 {
			sklog.Fatalf("Failed to split imageName: %v", parts)
		}
		imageNoTag := parts[0]
		imageRegex := regexp.MustCompile(fmt.Sprintf(`^(\s+image:\s+)(%s):.*$`, imageNoTag))

		// Loop over all the yaml files and update tags for the given imageName.
		for _, filename := range filenames {
			b, err := ioutil.ReadFile(filename)
			if err != nil {
				sklog.Errorf("Failed to read %q (skipping): %s", filename, err)
				continue
			}
			lines := strings.Split(string(b), "\n")
			for i, line := range lines {
				matches := imageRegex.FindStringSubmatch(line)
				if len(matches) != 3 {
					continue
				}
				changed[filename] = true
				lines[i] = matches[1] + image
			}
			if changed[filename] {
				err := util.WithWriteFile(filename, func(w io.Writer) error {
					_, err := w.Write([]byte(strings.Join(lines, "\n")))
					return err
				})
				if err != nil {
					sklog.Fatalf("Failed to write update config file %q: %s", filename, err)
				}
			}
		}
	}

	// Were any files updated?
	if len(changed) != 0 {
		filenameFlag := fmt.Sprintf("--filename=%s\n", strings.Join(changed.Keys(), ","))
		if !*dryRun {
			for filename, _ := range changed {
				msg, err := checkout.Git(ctx, "add", filepath.Base(filename))
				if err != nil {
					sklog.Fatalf("Failed to stage changes to the config repo: %s: %q", err, msg)
				}
			}
			msg, err := checkout.Git(ctx, "diff", "--name-only")
			if err != nil {
				sklog.Fatalf("Failed to diff :%s: %q", err, msg)
			}
			if msg == "" {
				sklog.Infof("Not pushing since no files changed.")
				return
			}
			msg, err = checkout.Git(ctx, "commit", "-m", *message)
			if err != nil {
				sklog.Fatalf("Failed to commit to the config repo: %s: %q", err, msg)
			}
			msg, err = checkout.Git(ctx, "push", "origin", "master")
			if err != nil {
				sklog.Fatalf("Failed to push the config repo: %s: %q", err, msg)
			}
			if err := exec.Run(context.Background(), &exec.Command{
				Name:      "kubectl",
				Args:      []string{"apply", filenameFlag},
				LogStderr: true,
				LogStdout: true,
			}); err != nil {
				sklog.Errorf("Failed to run: %s", err)
			}
		} else {
			fmt.Printf("\nkubectl apply %s\n", filenameFlag)
		}
	} else {
		fmt.Println("Nothing to do.")
	}
}
<commit_msg>pushk - Fix for detecting file changes.<commit_after>// pushk pushes a new version of an app.
//
// pushk
// pushk docserver
// pushk --rollback docserver
// pushk --project=skia-public docserver
// pushk --rollback --project=skia-public docserver
// pushk --dry-run my-new-service
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/gcr"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	repoUrlTemplate = "https://skia.googlesource.com/%s-config"
	repoBaseDir     = "/tmp"
	repoDirTemplate = "/tmp/%s-config"
)

func init() {
	flag.Usage = func() {
		fmt.Printf("Usage: pushk <flags> [zero or more image names]\n\n")
		fmt.Printf(`pushk pushes a new version of an app.

The command:
  1. Modifies the kubernetes yaml files with the new image.
  2. Commits the changes to the config repo.
  3. Applies the changes with kubectl.

The config is stored in a separate repo that will automaticaly be checked out
under /tmp.

The command applies the changes by default, or just changes the local yaml files
if --dry-run is supplied.

If no image names are supplied then pushk looks through all the yaml files for
appropriate images (ones that match the SERVER and project) and tries to push a
new image for each of them.

Examples:
  # Push the latest version of all images from the given container repository.
  pushk

  # Pusk an exact tag.
  pushk gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07

  # Push the latest version of docserver.
  pushk docserver --message="Fix bug #1234"

  # Push the latest version of docserver and iap-proxy
  pushk docserver iap-proxy

  # Rollback docserver.
  pushk --rollback docserver

  # Compute any changes a push to docserver will make, but do not apply them.
  pushk --dry-run docserver

`)
		flag.PrintDefaults()
	}
}

// toFullRepoURL converts the project name into a git repo URL.
func toFullRepoURL(s string) string {
	return fmt.Sprintf(repoUrlTemplate, s)

}

// toRepoDir converts the project name into a git repo directory name.
func toRepoDir(s string) string {
	return fmt.Sprintf(repoDirTemplate, s)
}

// flags
var (
	dryRun   = flag.Bool("dry-run", false, "If true then do not run the kubectl command to apply the changes, and do not commit the changes to the config repo.")
	message  = flag.String("message", "Push", "Message to go along with the change.")
	project  = flag.String("project", "skia-public", "The GCE project name.")
	rollback = flag.Bool("rollback", false, "If true go back to the second most recent image, otherwise use most recent image.")
)

var (
	validTag = regexp.MustCompile(`^\d\d\d\d-\d\d-\d\dT\d\d_\d\d_\d\dZ-.+$`)
)

// filter strips the list of tags down to only the ones that conform to our
// constraints and also checks that there are enough tags. The results
// are sorted in ascending order, so oldest tags are first, newest tags
// are last.
func filter(tags []string) ([]string, error) {
	validTags := []string{}
	for _, t := range tags {
		if validTag.MatchString(t) {
			validTags = append(validTags, t)
		}
	}
	sort.Strings(validTags)
	if len(validTags) == 0 {
		return nil, fmt.Errorf("Not enough tags returned.")
	}
	return validTags, nil
}

// findAllImageNames searches for all the images that come from the given
// project container registry across all the yaml files listed in filenames.
func findAllImageNames(filenames []string, server, project string) []string {
	// allImageRegex has the following groups returned on match:
	// 0 - the entire line
	// 1 - the image name
	allImageRegex := regexp.MustCompile(fmt.Sprintf(`(?m)^\s+image:\s+%s/%s/([^:]+):\S+\s*$`, server, project))
	filenameSet := util.StringSet{}
	for _, filename := range filenames {
		b, err := ioutil.ReadFile(filename)
		if err != nil {
			sklog.Errorf("Failed to read %q (skipping): %s", filename, err)
			continue
		}
		matches := allImageRegex.FindAllStringSubmatch(string(b), -1)
		for _, m := range matches {
			filenameSet[m[1]] = true
		}
	}
	return filenameSet.Keys()
}

// tagProvider is a type that returns the correct tag to push for the given imageName.
type tagProvider func(imageName string) ([]string, error)

// imageFromCmdLineImage handles image names, which can be either short, ala 'fiddler', or exact,
// such as gcr.io/skia-public/fiddler:694900e3ca9468784a5794dc53382d1c8411ab07, both of which can
// appear on the command-line.
func imageFromCmdLineImage(imageName string, tp tagProvider) (string, error) {
	if strings.HasPrefix(imageName, "gcr.io/") {
		if *rollback {
			return "", fmt.Errorf("Supplying a fully qualified image name and the --rollback flag are mutually exclusive.")
		}
		return imageName, nil
	}
	// Get all the tags for the selected image.
	tags, err := tp(imageName)
	if err != nil {
		return "", fmt.Errorf("Tag provider failed: %s", err)
	}

	// Filter the tags
	tags, err = filter(tags)
	if err != nil {
		return "", fmt.Errorf("Failed to filter: %s", err)
	}

	// Pick the target tag we want to move to.
	tag := tags[len(tags)-1]
	if *rollback {
		if len(tags) < 2 {
			return "", fmt.Errorf("No version to rollback to.")
		}
		tag = tags[len(tags)-2]
	}

	// The full docker image name and tag of the image we want to deploy.
	return fmt.Sprintf("%s/%s/%s:%s", gcr.SERVER, *project, imageName, tag), nil
}

func main() {
	common.Init()

	ctx := context.Background()
	repoDir := toRepoDir(*project)
	checkout, err := git.NewCheckout(ctx, toFullRepoURL(*project), repoBaseDir)
	if err != nil {
		sklog.Fatalf("Failed to check out config repo: %s", err)
	}
	if err := checkout.Update(ctx); err != nil {
		sklog.Fatalf("Failed to update repo: %s", err)
	}

	// Get all the yaml files.
	filenames, err := filepath.Glob(filepath.Join(repoDir, "*.yaml"))
	if err != nil {
		sklog.Fatal(err)
	}

	tokenSource := auth.NewGCloudTokenSource(*project)
	imageNames := flag.Args()
	if len(imageNames) == 0 {
		imageNames = findAllImageNames(filenames, gcr.SERVER, *project)
		if len(imageNames) == 0 {
			fmt.Printf("Failed to find any images that match kubernetes directory: %q and project: %q.", repoDir, *project)
			flag.Usage()
			os.Exit(1)
		}
	}
	sklog.Infof("Pushing the following images: %q", imageNames)

	gcrTagProvider := func(imageName string) ([]string, error) {
		return gcr.NewClient(tokenSource, *project, imageName).Tags()
	}

	changed := util.StringSet{}
	for _, imageName := range imageNames {
		image, err := imageFromCmdLineImage(imageName, gcrTagProvider)
		if err != nil {
			sklog.Fatal(err)
		}

		// imageRegex has the following groups returned on match:
		// 0 - the entire line
		// 1 - the prefix, i.e. image:, with correct spacing.
		// 2 - full image name
		// 3 - just the tag
		//
		// We pull out the 'prefix' so we can use it when
		// we rewrite the image: line so the indent level is
		// unchanged.
		parts := strings.SplitN(image, ":", 2)
		if len(parts) != 2 {
			sklog.Fatalf("Failed to split imageName: %v", parts)
		}
		imageNoTag := parts[0]
		imageRegex := regexp.MustCompile(fmt.Sprintf(`^(\s+image:\s+)(%s):.*$`, imageNoTag))

		// Loop over all the yaml files and update tags for the given imageName.
		for _, filename := range filenames {
			b, err := ioutil.ReadFile(filename)
			if err != nil {
				sklog.Errorf("Failed to read %q (skipping): %s", filename, err)
				continue
			}
			lines := strings.Split(string(b), "\n")
			for i, line := range lines {
				matches := imageRegex.FindStringSubmatch(line)
				if len(matches) != 3 {
					continue
				}
				changed[filename] = true
				lines[i] = matches[1] + image
			}
			if changed[filename] {
				err := util.WithWriteFile(filename, func(w io.Writer) error {
					_, err := w.Write([]byte(strings.Join(lines, "\n")))
					return err
				})
				if err != nil {
					sklog.Fatalf("Failed to write update config file %q: %s", filename, err)
				}
			}
		}
	}

	// Were any files updated?
	if len(changed) != 0 {
		filenameFlag := fmt.Sprintf("--filename=%s\n", strings.Join(changed.Keys(), ","))
		if !*dryRun {
			for filename, _ := range changed {
				msg, err := checkout.Git(ctx, "add", filepath.Base(filename))
				if err != nil {
					sklog.Fatalf("Failed to stage changes to the config repo: %s: %q", err, msg)
				}
			}
			msg, err := checkout.Git(ctx, "diff", "--cached", "--name-only")
			if err != nil {
				sklog.Fatalf("Failed to diff :%s: %q", err, msg)
			}
			if msg == "" {
				sklog.Infof("Not pushing since no files changed.")
				return
			}
			msg, err = checkout.Git(ctx, "commit", "-m", *message)
			if err != nil {
				sklog.Fatalf("Failed to commit to the config repo: %s: %q", err, msg)
			}
			msg, err = checkout.Git(ctx, "push", "origin", "master")
			if err != nil {
				sklog.Fatalf("Failed to push the config repo: %s: %q", err, msg)
			}
			if err := exec.Run(context.Background(), &exec.Command{
				Name:      "kubectl",
				Args:      []string{"apply", filenameFlag},
				LogStderr: true,
				LogStdout: true,
			}); err != nil {
				sklog.Errorf("Failed to run: %s", err)
			}
		} else {
			fmt.Printf("\nkubectl apply %s\n", filenameFlag)
		}
	} else {
		fmt.Println("Nothing to do.")
	}
}
<|endoftext|>