<commit_before>package dtable

import (
	"bytes"
	"github.com/fastfn/dendrite"
	"log"
	"sync"
)

// if node left, maintain consistency by finding local replicas and push them one step further if possible
// if node joined, find all keys in local tables that are < new_pred, copy them to new_pred and strip last replica for them
//                 for other keys, just copy them to all replicas as we might be in deficite
func (dt *DTable) Delegate(localVn, new_pred *dendrite.Vnode, event dendrite.RingEventType, mux sync.Mutex) {
	// get the handler for this vnode
	_, ok := dt.transport.GetVnodeHandler(localVn)
	if !ok {
		// can't do this
		return
	}

	//vn_table, _ := dt.table[localVn.String()]

	switch event {
	case dendrite.EvPredecessorLeft:
		dt.promote(localVn)
	case dendrite.EvPredecessorJoined:
		dt.demote(localVn, new_pred)
		/*
			log.Printf("Node joined me: %X  ... %X replicating to:\n", localVn.Id, new_pred.Id)
			for _, r := range replicas {
				log.Printf("\t - %X\n", r.Id)
			}
			// find all local keys that are < new predecessor
			for key_str, val := range vn_table {
				key, _ := hex.DecodeString(key_str)
				if dendrite.Between(key, localVn.Id, new_pred.Id, true) {
					// copy the key to new predecessor
					done_c := make(chan error)
					go dt.remoteSet(new_pred, key, val, 0, done_c)
					err := <-done_c
					if err != nil {
						log.Println("Dendrite::Delegate -- failed to delegate key to new predecessor:", err)
						continue
					}
					// remove the key from last replica unless last replica is our new predecessor
					if len(replicas) == 0 {
						continue
					}
					last_replica = replicas[len(replicas)-1]
					if last_replica.Host == new_pred.Host {
						continue
					}
					go dt.remoteSet(last_replica, key, nil, 0, done_c)
					err = <-done_c
					if err != nil {
						log.Println("Dendrite::Delegate - failed to strip key from last replica:", err)
					}
				} else {
					//for _, replica := range replicas {
					// err := dt.remoteSet(replica, key, val.Val)
					// if err != nil {
					// 	log.Println("Dendrite::Delegate -- failed to propagate key to replica:", err)
					// }
					//}
				}
			}
		*/
	default:
		return
	}
}

/*
type rvalue struct {
	Val           []byte
	timestamp     time.Time
	depth         int
	state         replicaState
	master        *dendrite.Vnode
	replicaVnodes []*dendrite.Vnode
}

type value struct {
	Val       []byte
	timestamp time.Time
	isReplica bool
	commited  bool
	rstate    replicaState
}
*/
// promote() - called when remote predecessor died or left the ring
// because we're only first REMOTE node from original master
// it doesn't mean that we're the actual successor for all the replicated data with depth 0
// if we are, we promote ourselves
// if not, we must find actual successor for each key, and promote that vnode for each key
func (dt *DTable) promote(vnode *dendrite.Vnode) {
	//log.Printf("Node left me: %X for %X now replicating to:\n", localVn.Id, new_pred.Id)
	rtable := dt.rtable[vnode.String()]
	vn_table := dt.table[vnode.String()]
	for key_str, ritem := range rtable {
		if ritem.replicaInfo.depth != 0 {
			continue
		}
		// check if we're real successor for this key
		succs, err := dt.ring.Lookup(1, ritem.keyHash)
		if err != nil {
			log.Printf("Could not promote key, Lookup() failed: %s\n", err.Error())
			continue
		}
		if bytes.Compare(succs[0].Id, vnode.Id) == 0 {
			// this key should be promoted locally
			new_ritem := ritem.dup()
			new_ritem.replicaInfo.vnodes[0] = nil
			new_ritem.commited = true
			vn_table[key_str] = new_ritem
			log.Printf("Promoted local key: %s - running replicator now replicas are %+v \n", key_str, new_ritem.replicaInfo.vnodes)
			delete(rtable, key_str)
			log.Printf("Promote calling replicateKey for key %s\n", key_str)
			dt.replicateKey(vnode, new_ritem, dt.ring.Replicas())
			log.Printf("Promote finishing key %s, replicaVnodes are: %+v\n", key_str, new_ritem.replicaInfo.vnodes)
		} else {
			// TODO promote remote vnode
			delete(rtable, key_str)
		}
	}
}

/* demote() - promotes new predecessor with keys from primary table
if new predecessor is local:
	- move all of my replica keys to new vnode
	- replica scheme of remote successors doesn't change here
	  we just need to update metadata on all replica nodes to reflect this change
if new predecessor is remote:
  - for all keys in primary table, that are <= new_pred.Id:
  	1. move key to demoted table and wait there for cleanup call from new master
  	2. call demoteKey() to commit to new_pred's primary table + let that vnode know where existing replicas are
  	3. demoteKey() will callback to cleanup each key from demoted table after it's written new replicas
  - handle replica-0 table such that:
  	1. for each key, check if master vnode is located on same physical node as new_pred
  	- if it is, we don't need to do anything because we're still natural remote successor
  	- if not
  		1. call demoteReplica() to let master know existing replica setup and about newRemoteSucc
  		2. master will reconfigure replicas around and delete unnecessary copies (if any)
*/
func (dt *DTable) demote(vnode, new_pred *dendrite.Vnode) {
	// determine if new_pred is on this node
	isLocal := false
	for _, lvn := range dt.ring.MyVnodes() {
		if lvn.Host == new_pred.Host {
			isLocal = true
		}
	}
	switch isLocal {
	case true:
		// move all replica keys to new vnode
		vn_rtable := dt.rtable[vnode.String()]
		for rkey, ritem := range vn_rtable {
			ritem.replicaInfo.vnodes[ritem.replicaInfo.depth] = new_pred
			dt.rtable[new_pred.String()][rkey] = ritem
			delete(vn_rtable, rkey)

			// update metadata on all replicas
			new_state := ritem.replicaInfo.state
			for idx, replica := range ritem.replicaInfo.vnodes {
				// skip ourselves
				if idx == ritem.replicaInfo.depth {
					continue
				}
				new_ritem := ritem.dup()
				new_ritem.replicaInfo.depth = idx
				new_ritem.replicaInfo.state = new_state

				err := dt.remoteSetReplicaInfo(replica, new_ritem)
				if err != nil {
					log.Println("Error updating replicaMeta on demote() -", err)
					new_state = replicaIncomplete
					continue
				}
			}
		}
	case false:
		// loop over primary table to find keys that should belong to new predecessor
		vn_table := dt.table[vnode.String()]
		for key_str, item := range vn_table {
			if dendrite.Between(vnode.Id, new_pred.Id, item.keyHash, true) {
				//log.Printf("Analyzed key for demoting %s and pushing to %s\n", key_str, new_pred.String())
				// copy the key to demoted table and remove it from primary one
				dt.demoted_table[vnode.String()][item.keyHashString()] = item.to_demoted(new_pred)
				delete(vn_table, key_str)
				done_c := make(chan error)
				go dt.remoteSet(vnode, new_pred, item, dt.ring.Replicas(), true, done_c)
				err := <-done_c
				if err != nil {
					log.Println("Error demoting key to new predecessor -", err)
					continue
				}
			}
		}
	}

}

// changeReplicas() -- callend when replica set changes
//
func (dt *DTable) changeReplicas(vnode *dendrite.Vnode, new_replicas []*dendrite.Vnode) {
	// loop over primary table and get replicaVnodes
	// loop over replicaVnodes and remove each
	// loop over new_replicas and push each
	// write metadata
	for _, item := range dt.table[vnode.String()] {
		dt.replicateKey(vnode, item, dt.ring.Replicas())
	}

	// write new replicas

	/*
		real_idx := 0
		success_replicas := make([]*dendrite.Vnode, dt.ring.Replicas())
		for _, remote := range new_replicas {
			if remote == nil {
				continue
			}
			log.Printf("replicating %s to: %x\n", key_str, remote.Id)
			new_ritem := item.dup()
			new_ritem.replicaInfo.state = replicaIncomplete
			new_ritem.commited = false
			new_ritem.replicaInfo.depth = real_idx

			if err := dt.remoteWriteReplica(vnode, remote, new_ritem); err != nil {
				log.Printf("changeReplicas() - error writing replica %s for key %s - %s\n", remote.String(), key_str, err.Error())
				continue
			}
			success_replicas[real_idx] = remote
			real_idx++
		}
		copy(item.replicaInfo.vnodes, success_replicas)

		// replicas have been written, lets now update metadata
		for idx, replica := range success_replicas {
			if replica == nil {
				continue
			}
			new_ritem := item.dup()
			new_ritem.replicaInfo.depth = idx
			new_ritem.replicaInfo.state = replicaStable
			err := dt.remoteSetReplicaInfo(replica, new_ritem)
			if err != nil {
				break
			}
		}
	*/

}

func (dt *DTable) replicateKey(vnode *dendrite.Vnode, reqItem *kvItem, limit int) {
	handler, _ := dt.transport.GetVnodeHandler(vnode)
	if handler == nil {
		return
	}
	// find remote successors to write replicas to
	remote_succs, err := handler.FindRemoteSuccessors(limit)
	if err != nil {
		return
	}

	// first, lets remove existing replicas
	for idx, existing := range reqItem.replicaInfo.vnodes {
		if existing == nil {
			continue
		}
		if err := dt.remoteClearReplica(existing, reqItem, false); err != nil {
			// lets add this replica to orphans
			reqItem.replicaInfo.orphan_vnodes = append(reqItem.replicaInfo.orphan_vnodes, existing)
			reqItem.replicaInfo.vnodes[idx] = nil
			continue
		} else {
			reqItem.replicaInfo.vnodes[idx] = nil
		}
	}

	// we set replicaState to stable if enough remote successors are found
	// otherwise, replicaState is still stable, but partial
	// if any of replica writes fail later on, we'll set the state to Incomplete
	var new_replica_state replicaState
	if len(remote_succs) >= dt.ring.Replicas() {
		new_replica_state = replicaStable
	} else {
		new_replica_state = replicaPartial
	}

	// now lets write replicas
	new_replicas := make([]*dendrite.Vnode, 0)

	for _, succ := range remote_succs {
		if succ == nil {
			continue
		}
		log.Printf("replicating to: %x\n", succ.Id)
		new_ritem := reqItem.dup()
		new_ritem.replicaInfo.state = replicaIncomplete
		new_ritem.commited = false

		err := dt.remoteWriteReplica(vnode, succ, new_ritem)
		if err != nil {
			log.Printf("Error writing replica to %s for key %s due to error: %s\n", succ.String(), new_ritem.keyHashString(), err.Error())
			new_replica_state = replicaIncomplete
			continue
		}
		new_replicas = append(new_replicas, succ)
	}

	// update metadata on original item
	reqItem.replicaInfo.vnodes = make([]*dendrite.Vnode, limit)
	for idx, new_replica := range new_replicas {
		reqItem.replicaInfo.vnodes[idx] = new_replica
	}
	reqItem.replicaInfo.state = new_replica_state

	// update metadata on successful replicas
	for idx, replica := range reqItem.replicaInfo.vnodes {
		if replica == nil {
			break
		}
		new_ritem := reqItem.dup()
		new_ritem.replicaInfo.depth = idx
		new_ritem.replicaInfo.state = new_replica_state

		err := dt.remoteSetReplicaInfo(replica, new_ritem)
		if err != nil {
			// this should not happen. It means another replica node failed in the meantime
			// need to trigger orphan cleaner, which will restart this process
			reqItem.replicaInfo.state = replicaIncomplete
			reqItem.replicaInfo.vnodes[idx] = nil
			reqItem.replicaInfo.orphan_vnodes = append(reqItem.replicaInfo.orphan_vnodes, replica)
			continue
		}
	}
}

/*
orphan cleaner after removing orhpans restarts replicateKey() if key's replica state is Incomplete
*/
<commit_msg>cleanup<commit_after>package dtable

import (
	"bytes"
	"github.com/fastfn/dendrite"
	"log"
	"sync"
)

*/
// promote() - called when remote predecessor died or left the ring
// because we're only first REMOTE node from original master
// it doesn't mean that we're the actual successor for all the replicated data with depth 0
// if we are, we promote ourselves
// if not, we must find actual successor for each key, and promote that vnode for each key
func (dt *DTable) promote(vnode *dendrite.Vnode) {
	//log.Printf("Node left me: %X for %X now replicating to:\n", localVn.Id, new_pred.Id)
	rtable := dt.rtable[vnode.String()]
	vn_table := dt.table[vnode.String()]
	for key_str, ritem := range rtable {
		if ritem.replicaInfo.depth != 0 {
			continue
		}
		// check if we're real successor for this key
		succs, err := dt.ring.Lookup(1, ritem.keyHash)
		if err != nil {
			log.Printf("Could not promote key, Lookup() failed: %s\n", err.Error())
			continue
		}
		if bytes.Compare(succs[0].Id, vnode.Id) == 0 {
			// this key should be promoted locally
			new_ritem := ritem.dup()
			new_ritem.replicaInfo.vnodes[0] = nil
			new_ritem.commited = true
			vn_table[key_str] = new_ritem
			log.Printf("Promoted local key: %s - running replicator now replicas are %+v \n", key_str, new_ritem.replicaInfo.vnodes)
			delete(rtable, key_str)
			log.Printf("Promote calling replicateKey for key %s\n", key_str)
			dt.replicateKey(vnode, new_ritem, dt.ring.Replicas())
			log.Printf("Promote finishing key %s, replicaVnodes are: %+v\n", key_str, new_ritem.replicaInfo.vnodes)
		} else {
			// TODO promote remote vnode
			delete(rtable, key_str)
		}
	}
}

/* demote() - promotes new predecessor with keys from primary table
if new predecessor is local:
	- move all of my replica keys to new vnode
	- replica scheme of remote successors doesn't change here
	  we just need to update metadata on all replica nodes to reflect this change
if new predecessor is remote:
  - for all keys in primary table, that are <= new_pred.Id:
  	1. move key to demoted table and wait there for cleanup call from new master
  	2. call demoteKey() to commit to new_pred's primary table + let that vnode know where existing replicas are
  	3. demoteKey() will callback to cleanup each key from demoted table after it's written new replicas
  - handle replica-0 table such that:
  	1. for each key, check if master vnode is located on same physical node as new_pred
  	- if it is, we don't need to do anything because we're still natural remote successor
  	- if not
  		1. call demoteReplica() to let master know existing replica setup and about newRemoteSucc
  		2. master will reconfigure replicas around and delete unnecessary copies (if any)
*/
func (dt *DTable) demote(vnode, new_pred *dendrite.Vnode) {
	// determine if new_pred is on this node
	isLocal := false
	for _, lvn := range dt.ring.MyVnodes() {
		if lvn.Host == new_pred.Host {
			isLocal = true
		}
	}
	switch isLocal {
	case true:
		// move all replica keys to new vnode
		vn_rtable := dt.rtable[vnode.String()]
		for rkey, ritem := range vn_rtable {
			ritem.replicaInfo.vnodes[ritem.replicaInfo.depth] = new_pred
			dt.rtable[new_pred.String()][rkey] = ritem
			delete(vn_rtable, rkey)

			// update metadata on all replicas
			new_state := ritem.replicaInfo.state
			for idx, replica := range ritem.replicaInfo.vnodes {
				// skip ourselves
				if idx == ritem.replicaInfo.depth {
					continue
				}
				new_ritem := ritem.dup()
				new_ritem.replicaInfo.depth = idx
				new_ritem.replicaInfo.state = new_state

				err := dt.remoteSetReplicaInfo(replica, new_ritem)
				if err != nil {
					log.Println("Error updating replicaMeta on demote() -", err)
					new_state = replicaIncomplete
					continue
				}
			}
		}
	case false:
		// loop over primary table to find keys that should belong to new predecessor
		vn_table := dt.table[vnode.String()]
		for key_str, item := range vn_table {
			if dendrite.Between(vnode.Id, new_pred.Id, item.keyHash, true) {
				//log.Printf("Analyzed key for demoting %s and pushing to %s\n", key_str, new_pred.String())
				// copy the key to demoted table and remove it from primary one
				dt.demoted_table[vnode.String()][item.keyHashString()] = item.to_demoted(new_pred)
				delete(vn_table, key_str)
				done_c := make(chan error)
				go dt.remoteSet(vnode, new_pred, item, dt.ring.Replicas(), true, done_c)
				err := <-done_c
				if err != nil {
					log.Println("Error demoting key to new predecessor -", err)
					continue
				}
			}
		}
	}

}

// changeReplicas() -- callend when replica set changes
//
func (dt *DTable) changeReplicas(vnode *dendrite.Vnode, new_replicas []*dendrite.Vnode) {
	for _, item := range dt.table[vnode.String()] {
		dt.replicateKey(vnode, item, dt.ring.Replicas())
	}
}

func (dt *DTable) replicateKey(vnode *dendrite.Vnode, reqItem *kvItem, limit int) {
	handler, _ := dt.transport.GetVnodeHandler(vnode)
	if handler == nil {
		return
	}
	// find remote successors to write replicas to
	remote_succs, err := handler.FindRemoteSuccessors(limit)
	if err != nil {
		return
	}

	// first, lets remove existing replicas
	for idx, existing := range reqItem.replicaInfo.vnodes {
		if existing == nil {
			continue
		}
		if err := dt.remoteClearReplica(existing, reqItem, false); err != nil {
			// lets add this replica to orphans
			reqItem.replicaInfo.orphan_vnodes = append(reqItem.replicaInfo.orphan_vnodes, existing)
			reqItem.replicaInfo.vnodes[idx] = nil
			continue
		} else {
			reqItem.replicaInfo.vnodes[idx] = nil
		}
	}

	// we set replicaState to stable if enough remote successors are found
	// otherwise, replicaState is still stable, but partial
	// if any of replica writes fail later on, we'll set the state to Incomplete
	var new_replica_state replicaState
	if len(remote_succs) >= dt.ring.Replicas() {
		new_replica_state = replicaStable
	} else {
		new_replica_state = replicaPartial
	}

	// now lets write replicas
	new_replicas := make([]*dendrite.Vnode, 0)

	for _, succ := range remote_succs {
		if succ == nil {
			continue
		}
		log.Printf("replicating to: %x\n", succ.Id)
		new_ritem := reqItem.dup()
		new_ritem.replicaInfo.state = replicaIncomplete
		new_ritem.commited = false

		err := dt.remoteWriteReplica(vnode, succ, new_ritem)
		if err != nil {
			log.Printf("Error writing replica to %s for key %s due to error: %s\n", succ.String(), new_ritem.keyHashString(), err.Error())
			new_replica_state = replicaIncomplete
			continue
		}
		new_replicas = append(new_replicas, succ)
	}

	// update metadata on original item
	reqItem.replicaInfo.vnodes = make([]*dendrite.Vnode, limit)
	for idx, new_replica := range new_replicas {
		reqItem.replicaInfo.vnodes[idx] = new_replica
	}
	reqItem.replicaInfo.state = new_replica_state

	// update metadata on successful replicas
	for idx, replica := range reqItem.replicaInfo.vnodes {
		if replica == nil {
			break
		}
		new_ritem := reqItem.dup()
		new_ritem.replicaInfo.depth = idx
		new_ritem.replicaInfo.state = new_replica_state

		err := dt.remoteSetReplicaInfo(replica, new_ritem)
		if err != nil {
			// this should not happen. It means another replica node failed in the meantime
			// need to trigger orphan cleaner, which will restart this process
			reqItem.replicaInfo.state = replicaIncomplete
			reqItem.replicaInfo.vnodes[idx] = nil
			reqItem.replicaInfo.orphan_vnodes = append(reqItem.replicaInfo.orphan_vnodes, replica)
			continue
		}
	}
}

/*
orphan cleaner after removing orhpans restarts replicateKey() if key's replica state is Incomplete
*/
<|endoftext|>


<commit_before>package vkv

import (
	"reflect"
	"sync"
	"testing"
	"time"
)

func check(e error) {
	if e != nil {
		panic(e)
	}
}
func TestDB(t *testing.T) {
	db, err := New("db_base")
	defer db.Destroy()
	if err != nil {
		t.Fatalf("Error creating db %v", err)
	}
	// testing low level binary encoded uint64
	valUint, err := db.getUint64([]byte("foo2"))
	check(err)
	if valUint != uint64(0) {
		t.Errorf("Uninitialized uint64 key should return 0")
	}
	err = db.putUint64([]byte("foo2"), uint64(5))
	check(err)
	valUint, err = db.getUint64([]byte("foo2"))
	check(err)
	if valUint != uint64(5) {
		t.Error("Key should be set to 5")
	}
	err = db.incrUint64([]byte("foo2"), 2)
	check(err)
	valUint, err = db.getUint64([]byte("foo2"))
	check(err)
	if valUint != uint64(7) {
		t.Error("Key should be set to 7")
	}
	var wg sync.WaitGroup
	// Test the mutex
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err = db.incrUint64([]byte("foo3"), 1)
			check(err)
		}()
	}
	wg.Wait()
	valUint, err = db.getUint64([]byte("foo3"))
	check(err)
	if valUint != uint64(10) {
		t.Errorf("Key foo3 should be set to 10, got %v", valUint)
	}
	keys, err := db.Keys("", "\xff", 0)
	check(err)
	if len(keys) != 0 {
		t.Errorf("Keys() should be empty: %q", keys)
	}
	res, err := db.Put("test_key_1", "test_value_1", -1)
	check(err)
	res2, err := db.Get("test_key_1", -1)
	if !reflect.DeepEqual(res, res2) {
		t.Errorf("bad KeyValue result got %+v, expected %+v", res2, res)
	}
	res, err = db.Put("test_key_1", "test_value_1.1", -1)
	check(err)
	res2, err = db.Get("test_key_1", -1)
	if !reflect.DeepEqual(res, res2) {
		t.Errorf("bad KeyValue result got %+v, expected %+v", res2, res)
	}
	versions, err := db.Versions("test_key_1", 0, int(time.Now().UTC().UnixNano()), 0)
	check(err)
	t.Logf("v:%q", versions)
}
<commit_msg>vkv: more tests<commit_after>package vkv

import (
	"reflect"
	"sync"
	"testing"
	"time"
)

func check(e error) {
	if e != nil {
		panic(e)
	}
}
func TestDB(t *testing.T) {
	db, err := New("db_base")
	defer db.Destroy()
	if err != nil {
		t.Fatalf("Error creating db %v", err)
	}
	// testing low level binary encoded uint64
	valUint, err := db.getUint64([]byte("foo2"))
	check(err)
	if valUint != uint64(0) {
		t.Errorf("Uninitialized uint64 key should return 0")
	}
	err = db.putUint64([]byte("foo2"), uint64(5))
	check(err)
	valUint, err = db.getUint64([]byte("foo2"))
	check(err)
	if valUint != uint64(5) {
		t.Error("Key should be set to 5")
	}
	err = db.incrUint64([]byte("foo2"), 2)
	check(err)
	valUint, err = db.getUint64([]byte("foo2"))
	check(err)
	if valUint != uint64(7) {
		t.Error("Key should be set to 7")
	}
	var wg sync.WaitGroup
	// Test the mutex
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			err = db.incrUint64([]byte("foo3"), 1)
			check(err)
		}()
	}
	wg.Wait()
	valUint, err = db.getUint64([]byte("foo3"))
	check(err)
	if valUint != uint64(10) {
		t.Errorf("Key foo3 should be set to 10, got %v", valUint)
	}
	keys, err := db.Keys("", "\xff", 0)
	check(err)
	if len(keys) != 0 {
		t.Errorf("Keys() should be empty: %q", keys)
	}
	res, err := db.Put("test_key_1", "test_value_1", -1)
	check(err)
	res2, err := db.Get("test_key_1", -1)
	if !reflect.DeepEqual(res, res2) {
		t.Errorf("bad KeyValue result got %+v, expected %+v", res2, res)
	}
	res, err = db.Put("test_key_1", "test_value_1.1", -1)
	check(err)
	res2, err = db.Get("test_key_1", -1)
	if !reflect.DeepEqual(res, res2) {
		t.Errorf("bad KeyValue result got %+v, expected %+v", res2, res)
	}
	versions, err := db.Versions("test_key_1", 0, int(time.Now().UTC().UnixNano()), 0)
	check(err)
	if len(versions.Versions) != 2 {
		t.Errorf("key test_key_1 should have 2 versions, got %d", len(versions.Versions))
	}
	if versions.Versions[1].Value != res2.Value {
		t.Errorf("bad KeyValue result got %+v, expected %+v", versions.Versions[0].Value, res2.Value)
	}
	keys, err = db.Keys("", "\xff", 0)
	check(err)
	if len(keys) != 1 || keys[0] != res2.Key {
		t.Errorf("bad Keys() result, expected [%v], got %q", res2.Key, keys)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/alexflint/go-arg"
	"github.com/biogo/hts/bgzf"
	"github.com/brentp/xopen"
	ggd_utils "github.com/gogetdata/ggd-utils"
)

var args struct {
	Path   string `arg:"positional"`
	Genome string `arg:"-g,required,help:a genome file of chromosome sizes and order"`
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var VCFCommentAfterHeader = errors.New("comment line after non-header #CHROM line in VCF")

type chromStartGetter func([]byte) ([]byte, int, error)

var get_vcf_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	cpr := bytes.SplitN(line, []byte{'\t'}, 3)
	chrom := cpr[0]
	pos, err := strconv.Atoi(string(cpr[1]))
	return chrom, pos, err
}

var get_gff_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	toks := bytes.SplitN(line, []byte{'\t'}, 5)
	start, err := strconv.Atoi(string(toks[3]))
	return toks[0], start, err
}

func main() {

	p := arg.MustParse(&args)
	if args.Path == "" || args.Genome == "" {
		p.Fail("must specify a path to check and a genome file")
	}
	gf, err := ggd_utils.ReadGenomeFile(args.Genome, "")
	if err != nil {
		log.Fatal(err)
	}

	if strings.HasSuffix(args.Path, ".vcf.gz") {
		checkVCF(args.Path, gf)
	} else if strings.HasSuffix(args.Path, ".bed.gz") {
		checkTab(args.Path, gf, get_vcf_chrom_start)
	} else {
		found := false
		for _, suff := range []string{"gff", "gtf", "gff3", "gff2"} {
			if strings.HasSuffix(args.Path, suff) || strings.HasSuffix(args.Path, suff+".gz") {
				found = true
				checkTab(args.Path, gf, get_gff_chrom_start)
				break
			}
		}
		if !found {
			log.Fatalf("Don't know how to check this type of file: %s\n", args.Path)
		}
	}
}

func checkLine(iline int, line []byte, lastChrom []byte, lastStart int,
	get_chrom_start func([]byte) ([]byte, int, error),
	cmp func(a, b string) bool) ([]byte, int, error) {

	chrom, start, err := get_chrom_start(line)
	if err != nil {
		return chrom, start, err
	}
	if !bytes.Equal(chrom, lastChrom) {
		if len(lastChrom) != 0 {
			if !cmp(string(lastChrom), string(chrom)) {
				return chrom, start, fmt.Errorf("chromosomes not in specified sort order: %s, %s at line %d\nuse gsort (https://github.com/brentp/gsort/) to order according to the genome file", lastChrom, chrom, iline)
			}
		}
		lastChrom = chrom
		lastStart = start
	} else {
		if start < lastStart {
			return chrom, start, fmt.Errorf("positions not sorted: %d => %d at line %d\n", lastStart, start, iline)
		}
		lastStart = start
	}
	if start < 0 {
		return chrom, start, fmt.Errorf("negative position at line: %d (%d)", iline, start)
	}

	return chrom, start, nil
}

func checkTab(path string, gf *ggd_utils.GenomeFile, getter chromStartGetter) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatalf("BED: %s should have a .tbi", path)
	}
	fh, err := os.Open(path)
	if err != nil {
		log.Fatalf("BED: unable to open file: %s", path)
	}
	if ok, _ := bgzf.HasEOF(fh); !ok {
		log.Fatal("missing EOF")
	}

	bgz, err := bgzf.NewReader(fh, 1)
	defer bgz.Close()
	defer fh.Close()

	rdr := bufio.NewReader(bgz)
	iline := 1
	check(err)
	afterHeader := false
	lastChrom := []byte("")
	lastPos := -1
	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if len(line) == 0 {
			continue
		}
		if line[0] == '#' || bytes.HasPrefix(line, []byte("track")) || bytes.HasPrefix(line, []byte("browser")) {
			if afterHeader {
				log.Fatalf("found comment/header line after header at line: %d", iline)
			}
		} else {
			afterHeader = true
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, getter, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}

func checkVCF(path string, gf *ggd_utils.GenomeFile) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatal("VCF should have a .tbi or .csi")
	}

	fh, err := os.Open(path)
	check(err)
	defer fh.Close()
	if ok, _ := bgzf.HasEOF(fh); !ok {
		log.Fatalf("missing EOF in %s", path)
	}

	bgz, err := bgzf.NewReader(fh, 1)
	check(err)
	defer bgz.Close()
	rdr := bufio.NewReader(bgz)

	afterHeader := false
	iline := 1
	lastChrom := []byte("")
	lastPos := -1

	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if iline == 1 {
			if !bytes.HasPrefix(line, []byte("##fileformat=VCF")) {
				log.Println(string(line))
				log.Fatal("VCF header line '##fileformat=VCF... not found")
			}
		}
		if line[0] == '#' {
			if bytes.HasPrefix(line, []byte("#CHROM\t")) {
				afterHeader = true
			} else if afterHeader {
				log.Fatal(VCFCommentAfterHeader, iline)
			}
		} else {
			if !afterHeader {
				log.Fatal("VCF header line '#CHROM\t ... not found")
			}
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, get_vcf_chrom_start, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}
<commit_msg>Add support for bedGraph.gz files (#1)<commit_after>package main

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/alexflint/go-arg"
	"github.com/biogo/hts/bgzf"
	"github.com/brentp/xopen"
	ggd_utils "github.com/gogetdata/ggd-utils"
)

var args struct {
	Path   string `arg:"positional"`
	Genome string `arg:"-g,required,help:a genome file of chromosome sizes and order"`
}

func check(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

var VCFCommentAfterHeader = errors.New("comment line after non-header #CHROM line in VCF")

type chromStartGetter func([]byte) ([]byte, int, error)

var get_vcf_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	cpr := bytes.SplitN(line, []byte{'\t'}, 3)
	chrom := cpr[0]
	pos, err := strconv.Atoi(string(cpr[1]))
	return chrom, pos, err
}

var get_gff_chrom_start chromStartGetter = func(line []byte) ([]byte, int, error) {
	toks := bytes.SplitN(line, []byte{'\t'}, 5)
	start, err := strconv.Atoi(string(toks[3]))
	return toks[0], start, err
}

func main() {

	p := arg.MustParse(&args)
	if args.Path == "" || args.Genome == "" {
		p.Fail("must specify a path to check and a genome file")
	}
	gf, err := ggd_utils.ReadGenomeFile(args.Genome, "")
	if err != nil {
		log.Fatal(err)
	}

	if strings.HasSuffix(args.Path, ".vcf.gz") {
		checkVCF(args.Path, gf)
	} else if strings.HasSuffix(args.Path, ".bed.gz") || strings.HasSuffix(strings.ToLower(args.Path), ".bedgraph.gz") {
		checkTab(args.Path, gf, get_vcf_chrom_start)
	} else {
		found := false
		for _, suff := range []string{"gff", "gtf", "gff3", "gff2"} {
			if strings.HasSuffix(args.Path, suff) || strings.HasSuffix(args.Path, suff+".gz") {
				found = true
				checkTab(args.Path, gf, get_gff_chrom_start)
				break
			}
		}
		if !found {
			log.Fatalf("Don't know how to check this type of file: %s\n", args.Path)
		}
	}
}

func checkLine(iline int, line []byte, lastChrom []byte, lastStart int,
	get_chrom_start func([]byte) ([]byte, int, error),
	cmp func(a, b string) bool) ([]byte, int, error) {

	chrom, start, err := get_chrom_start(line)
	if err != nil {
		return chrom, start, err
	}
	if !bytes.Equal(chrom, lastChrom) {
		if len(lastChrom) != 0 {
			if !cmp(string(lastChrom), string(chrom)) {
				return chrom, start, fmt.Errorf("chromosomes not in specified sort order: %s, %s at line %d\nuse gsort (https://github.com/brentp/gsort/) to order according to the genome file", lastChrom, chrom, iline)
			}
		}
		lastChrom = chrom
		lastStart = start
	} else {
		if start < lastStart {
			return chrom, start, fmt.Errorf("positions not sorted: %d => %d at line %d\n", lastStart, start, iline)
		}
		lastStart = start
	}
	if start < 0 {
		return chrom, start, fmt.Errorf("negative position at line: %d (%d)", iline, start)
	}

	return chrom, start, nil
}

func checkTab(path string, gf *ggd_utils.GenomeFile, getter chromStartGetter) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatalf("BED: %s should have a .tbi", path)
	}
	fh, err := os.Open(path)
	if err != nil {
		log.Fatalf("BED: unable to open file: %s", path)
	}
	if ok, _ := bgzf.HasEOF(fh); !ok {
		log.Fatal("missing EOF")
	}

	bgz, err := bgzf.NewReader(fh, 1)
	defer bgz.Close()
	defer fh.Close()

	rdr := bufio.NewReader(bgz)
	iline := 1
	check(err)
	afterHeader := false
	lastChrom := []byte("")
	lastPos := -1
	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if len(line) == 0 {
			continue
		}
		if line[0] == '#' || bytes.HasPrefix(line, []byte("track")) || bytes.HasPrefix(line, []byte("browser")) {
			if afterHeader {
				log.Fatalf("found comment/header line after header at line: %d", iline)
			}
		} else {
			afterHeader = true
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, getter, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}

func checkVCF(path string, gf *ggd_utils.GenomeFile) {
	if !(xopen.Exists(path+".tbi") || xopen.Exists(path+".csi")) {
		log.Fatal("VCF should have a .tbi or .csi")
	}

	fh, err := os.Open(path)
	check(err)
	defer fh.Close()
	if ok, _ := bgzf.HasEOF(fh); !ok {
		log.Fatalf("missing EOF in %s", path)
	}

	bgz, err := bgzf.NewReader(fh, 1)
	check(err)
	defer bgz.Close()
	rdr := bufio.NewReader(bgz)

	afterHeader := false
	iline := 1
	lastChrom := []byte("")
	lastPos := -1

	for {
		line, err := rdr.ReadBytes('\n')
		if err == io.EOF {
			break
		}
		if iline == 1 {
			if !bytes.HasPrefix(line, []byte("##fileformat=VCF")) {
				log.Println(string(line))
				log.Fatal("VCF header line '##fileformat=VCF... not found")
			}
		}
		if line[0] == '#' {
			if bytes.HasPrefix(line, []byte("#CHROM\t")) {
				afterHeader = true
			} else if afterHeader {
				log.Fatal(VCFCommentAfterHeader, iline)
			}
		} else {
			if !afterHeader {
				log.Fatal("VCF header line '#CHROM\t ... not found")
			}
			lastChrom, lastPos, err = checkLine(iline, line, lastChrom, lastPos, get_vcf_chrom_start, gf.Less)
			if err != nil {
				log.Fatal(err)
			}
			if gf.Lengths[string(lastChrom)] < lastPos {
				log.Fatalf("position: %d beyond end of chromosome %s", lastPos, lastChrom)
			}
		}
		iline += 1
	}
}
<|endoftext|>


<commit_before>package app

import (
	"context"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	alicloud "k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager"
	"k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager/utils/metric"
	componentbaseconfig "k8s.io/component-base/config"
	"k8s.io/klog"
	ccfg "k8s.io/kubernetes/cmd/cloud-controller-manager/app/apis/config"
	genericcontrollermanager "k8s.io/kubernetes/cmd/controller-manager/app"
	kubectrlmgrconfig "k8s.io/kubernetes/pkg/controller/apis/config"
	serviceconfig "k8s.io/kubernetes/pkg/controller/service/config"
	"k8s.io/kubernetes/pkg/master/ports"
	"runtime"
	"time"

	"fmt"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apiserver/pkg/server/healthz"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	v1c "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/leaderelection"
	"k8s.io/client-go/tools/leaderelection/resourcelock"
	"k8s.io/client-go/tools/record"
	"k8s.io/cloud-provider"
	"k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager/controller/route"
	"k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager/controller/service"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/controller"
	"k8s.io/kubernetes/pkg/util/configz"
	"math/rand"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"strconv"
)

const (
	// ControllerStartJitter is the jitter value used when starting controller managers.
	ControllerStartJitter = 1.0
)

// ServerCCM is the main context object for the controller manager.
type ServerCCM struct {
	ccfg.CloudControllerManagerConfiguration
	restConfig *rest.Config
	election   *kubernetes.Clientset
	client     *kubernetes.Clientset
	recorder   record.EventRecorder
	Master     string
	Kubeconfig string
	cloud      cloudprovider.Interface

	// NodeStatusUpdateFrequency is the frequency at which the controller
	// updates nodes' status
	NodeStatusUpdateFrequency metav1.Duration
}

// NewServerCCM creates a new ExternalCMServer with a default config.
func NewServerCCM() *ServerCCM {
	ccm := ServerCCM{
		// Part of these default values also present in 'cmd/kube-controller-manager/app/options/options.go'.
		// Please keep them in sync when doing update.
		CloudControllerManagerConfiguration: ccfg.CloudControllerManagerConfiguration{
			Generic: kubectrlmgrconfig.GenericControllerManagerConfiguration{
				Port:            ports.CloudControllerManagerPort,
				Address:         "0.0.0.0",
				MinResyncPeriod: metav1.Duration{Duration: 5 * time.Minute},
				ClientConnection: componentbaseconfig.ClientConnectionConfiguration{
					ContentType: "application/vnd.kubernetes.protobuf",
					QPS:         20.0,
					Burst:       30,
				},
				LeaderElection: componentbaseconfig.LeaderElectionConfiguration{
					LeaderElect:   false,
					LeaseDuration: metav1.Duration{Duration: 15 * time.Second},
					RenewDeadline: metav1.Duration{Duration: 10 * time.Second},
					RetryPeriod:   metav1.Duration{Duration: 2 * time.Second},
					ResourceLock:  resourcelock.LeasesResourceLock,
				},
				ControllerStartInterval: metav1.Duration{Duration: 0 * time.Second},
			},
			KubeCloudShared: kubectrlmgrconfig.KubeCloudSharedConfiguration{
				NodeMonitorPeriod:         metav1.Duration{Duration: 5 * time.Second},
				ClusterName:               "kubernetes",
				ConfigureCloudRoutes:      true,
				RouteReconciliationPeriod: metav1.Duration{Duration: 10 * time.Second},
			},
			ServiceController: serviceconfig.ServiceControllerConfiguration{
				ConcurrentServiceSyncs: 3,
			},
		},
		NodeStatusUpdateFrequency: metav1.Duration{Duration: 5 * time.Minute},
	}
	ccm.Generic.LeaderElection.LeaderElect = true
	return &ccm
}

func createRecorder(client *kubernetes.Clientset) record.EventRecorder {
	cast := record.NewBroadcaster()
	cast.StartLogging(klog.Infof)
	cast.StartRecordingToSink(
		&v1c.EventSinkImpl{
			Interface: v1c.New(client.CoreV1().RESTClient()).Events(""),
		},
	)
	return cast.NewRecorder(
		legacyscheme.Scheme,
		v1.EventSource{Component: "cloud-controller-manager"},
	)
}

func client(ccm *ServerCCM) (*kubernetes.Clientset, *kubernetes.Clientset, *rest.Config, error) {
	kubeconfig, err := clientcmd.BuildConfigFromFlags(ccm.Master, ccm.Kubeconfig)
	if err != nil {
		return nil, nil, kubeconfig, err
	}

	kubeconfig.QPS = ccm.Generic.ClientConnection.QPS
	kubeconfig.Burst = int(ccm.Generic.ClientConnection.Burst)
	kubeconfig.ContentConfig.ContentType = ccm.Generic.ClientConnection.ContentType

	client, err := kubernetes.NewForConfig(
		rest.AddUserAgent(
			kubeconfig,
			"cloud-controller-manager",
		),
	)
	if err != nil {
		klog.Fatalf("invalid API configuration: %v", err)
	}
	leader := kubernetes.NewForConfigOrDie(
		rest.AddUserAgent(
			kubeconfig,
			"leader-election",
		),
	)
	return client, leader, kubeconfig, nil
}

func (ccm *ServerCCM) initialization() error {
	if ccm.KubeCloudShared.CloudProvider.Name == "" {
		return fmt.Errorf("--cloud-provider cannot be empty")
	}

	alicloud.CloudConfigFile = ccm.KubeCloudShared.CloudProvider.CloudConfigFile
	cloud, err := cloudprovider.InitCloudProvider(
		ccm.KubeCloudShared.CloudProvider.Name,
		ccm.KubeCloudShared.CloudProvider.CloudConfigFile,
	)
	if err != nil {
		return fmt.Errorf("cloud provider could not be initialized: %v", err)
	}

	ccm.cloud = cloud
	if !cloud.HasClusterID() {
		if ccm.KubeCloudShared.AllowUntaggedCloud {
			klog.Warning("detected a cluster without a ClusterID.  A ClusterID will " +
				"be required in the future.  Please tag your cluster to avoid any future issues")
		} else {
			return fmt.Errorf("no ClusterID found.  A ClusterID is required for the " +
				"cloud provider to function properly.  This check can be bypassed by " +
				"setting the allow-untagged-cloud option")
		}
	}

	cfg, err := configz.New("componentconfig")
	if err != nil {
		return fmt.Errorf("unable to register configz: %v", err)
	}
	cfg.Set(ccm.CloudControllerManagerConfiguration)

	ccm.client, ccm.election, ccm.restConfig, err = client(ccm)
	if err != nil {
		return fmt.Errorf("create client error: %s", err.Error())
	}
	ccm.recorder = createRecorder(ccm.client)
	return err
}

func (ccm *ServerCCM) Start() error {
	if err := ccm.initialization(); err != nil {
		return fmt.Errorf("verify ccm config: %s", err.Error())
	}

	// Start the external controller manager server
	go func() {
		mux := http.NewServeMux()
		healthz.InstallHandler(mux)
		if ccm.Generic.Debugging.EnableProfiling {
			mux.HandleFunc("/debug/pprof/", pprof.Index)
			mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
			mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
			mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
			if ccm.Generic.Debugging.EnableContentionProfiling {
				runtime.SetBlockProfileRate(1)
			}
		}
		configz.InstallHandler(mux)
		metric.RegisterPrometheus()
		mux.Handle("/metrics", promhttp.Handler())
		server := &http.Server{
			Addr:    net.JoinHostPort(ccm.Generic.Address, strconv.Itoa(int(ccm.Generic.Port))),
			Handler: mux,
		}
		klog.Fatal(server.ListenAndServe())
	}()
	return nil
}

func (ccm *ServerCCM) MainLoop(ctx context.Context) {
	var (
		builder controller.ControllerClientBuilder
	)
	builder = controller.SimpleControllerClientBuilder{
		ClientConfig: ccm.restConfig,
	}
	if ccm.KubeCloudShared.UseServiceAccountCredentials {
		builder = controller.SAControllerClientBuilder{
			ClientConfig:         rest.AnonymousClientConfig(ccm.restConfig),
			CoreClient:           ccm.client.CoreV1(),
			AuthenticationClient: ccm.client.AuthenticationV1(),
			Namespace:            "kube-system",
		}
	}
	panic(fmt.Sprintf("unreachable: %v", RunControllers(ccm, builder, nil)))
}

// Run runs the ExternalCMServer.  This should never exit.
func Run(ccm *ServerCCM) error {
	if err := ccm.Start(); err != nil {
		return err
	}

	route.Options = route.RoutesOptions{
		ClusterCIDR:               ccm.KubeCloudShared.ClusterCIDR,
		AllocateNodeCIDRs:         ccm.KubeCloudShared.AllocateNodeCIDRs,
		ConfigCloudRoutes:         ccm.KubeCloudShared.ConfigureCloudRoutes,
		RouteReconciliationPeriod: ccm.KubeCloudShared.RouteReconciliationPeriod,
		ControllerStartInterval:   ccm.Generic.ControllerStartInterval,
	}

	if !ccm.Generic.LeaderElection.LeaderElect {
		ccm.MainLoop(context.TODO())
	}

	// Identity used to distinguish between multiple cloud controller manager instances
	id, err := os.Hostname()
	if err != nil {
		return fmt.Errorf("hostname: %s", err.Error())
	}

	// Lock required for leader election
	rl, err := resourcelock.New(
		ccm.Generic.LeaderElection.ResourceLock,
		"kube-system",
		"ccm",
		ccm.election.CoreV1(),
		ccm.election.CoordinationV1(),
		resourcelock.ResourceLockConfig{
			Identity:      "ccm-" + id,
			EventRecorder: ccm.recorder,
		})
	if err != nil {
		return fmt.Errorf("creating etcd resource lock: %v", err)
	}

	// Try and become the leader and start cloud controller manager loops
	leaderelection.RunOrDie(
		context.Background(),
		leaderelection.LeaderElectionConfig{
			Lock: rl,
			Callbacks: leaderelection.LeaderCallbacks{
				OnStartedLeading: ccm.MainLoop,
				OnStoppedLeading: func() {
					klog.Fatalf("leaderelection lost")
				},
			},
			LeaseDuration: ccm.Generic.LeaderElection.LeaseDuration.Duration,
			RenewDeadline: ccm.Generic.LeaderElection.RenewDeadline.Duration,
			RetryPeriod:   ccm.Generic.LeaderElection.RetryPeriod.Duration,
		},
	)
	panic("unreachable")
}

// RunControllers starts the cloud specific controller loops.
func RunControllers(
	ccm *ServerCCM,
	clientBuilder controller.ControllerClientBuilder,
	stop <-chan struct{},
) error {

	if ccm.cloud != nil {
		// Initialize the cloud provider with a reference to the clientBuilder
		ccm.cloud.Initialize(clientBuilder, stop)
	}
	client := clientBuilder.ClientOrDie("shared-informers")

	ifactory := informers.NewSharedInformerFactory(client, resyncPeriod(ccm)())

	//if err := runControllerPV(ccm, clientBuilder, stop); err != nil {
	//	return fmt.Errorf("run pvcontroller: %s", err.Error())
	//}
	time.Sleep(wait.Jitter(ccm.Generic.ControllerStartInterval.Duration, ControllerStartJitter))

	if err := runControllerService(ccm, clientBuilder, ifactory, stop); err != nil {
		return fmt.Errorf("run service controller: %s", err.Error())
	}

	time.Sleep(wait.Jitter(ccm.Generic.ControllerStartInterval.Duration, ControllerStartJitter))

	// If apiserver is not running we should wait for some time and fail
	// only then. This is particularly important when we start apiserver
	// and controller manager at the same time.
	err := genericcontrollermanager.WaitForAPIServer(ccm.client, 10*time.Second)
	if err != nil {
		return fmt.Errorf("failed to get api versions from server: %v", err)
	}

	ifactory.Start(stop)
	klog.Infof("informer started")

	select {}
}

//func runControllerPV(
//	ccm *ServerCCM,
//	builder controller.ControllerClientBuilder,
//	stop <-chan struct{},
//) error {
//
//	con := ccontroller.NewPersistentVolumeLabelController(
//		builder.ClientOrDie("pvl-controller"),
//		ccm.cloud,
//	)
//	go con.Run(5, stop)
//	return nil
//}

func runControllerService(
	ccm *ServerCCM,
	builder controller.ControllerClientBuilder,
	informer informers.SharedInformerFactory,
	stop <-chan struct{},
) error {
	cloudslb, ok := ccm.cloud.LoadBalancer()
	if !ok {
		return fmt.Errorf("loadbalancer interface must be implemented")
	}

	scon, err := service.NewController(
		cloudslb,
		builder.ClientOrDie("cloud-controller-manager"),
		informer,
		ccm.KubeCloudShared.ClusterName,
	)
	if err != nil {
		return fmt.Errorf("failed to start service controller: %v", err)
	}
	go scon.Run(stop, int(ccm.ServiceController.ConcurrentServiceSyncs))
	return nil
}

func resyncPeriod(ccm *ServerCCM) func() time.Duration {
	return func() time.Duration {
		factor := rand.Float64() + 1
		return time.Duration(float64(ccm.Generic.MinResyncPeriod.Nanoseconds()) * factor)
	}
}
<commit_msg>set different ResourceLock addcording to kubernetes version<commit_after>package app

import (
	"context"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/version"
	alicloud "k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager"
	"k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager/utils/metric"
	componentbaseconfig "k8s.io/component-base/config"
	"k8s.io/klog"
	ccfg "k8s.io/kubernetes/cmd/cloud-controller-manager/app/apis/config"
	genericcontrollermanager "k8s.io/kubernetes/cmd/controller-manager/app"
	kubectrlmgrconfig "k8s.io/kubernetes/pkg/controller/apis/config"
	serviceconfig "k8s.io/kubernetes/pkg/controller/service/config"
	"k8s.io/kubernetes/pkg/master/ports"
	"runtime"
	"time"

	"fmt"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apiserver/pkg/server/healthz"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	v1c "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/leaderelection"
	"k8s.io/client-go/tools/leaderelection/resourcelock"
	"k8s.io/client-go/tools/record"
	"k8s.io/cloud-provider"
	"k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager/controller/route"
	"k8s.io/cloud-provider-alibaba-cloud/cloud-controller-manager/controller/service"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/controller"
	"k8s.io/kubernetes/pkg/util/configz"
	"math/rand"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"strconv"
)

const (
	// ControllerStartJitter is the jitter value used when starting controller managers.
	ControllerStartJitter = 1.0
)

// ServerCCM is the main context object for the controller manager.
type ServerCCM struct {
	ccfg.CloudControllerManagerConfiguration
	restConfig *rest.Config
	election   *kubernetes.Clientset
	client     *kubernetes.Clientset
	recorder   record.EventRecorder
	Master     string
	Kubeconfig string
	cloud      cloudprovider.Interface

	// NodeStatusUpdateFrequency is the frequency at which the controller
	// updates nodes' status
	NodeStatusUpdateFrequency metav1.Duration
}

// NewServerCCM creates a new ExternalCMServer with a default config.
func NewServerCCM() *ServerCCM {
	ccm := ServerCCM{
		// Part of these default values also present in 'cmd/kube-controller-manager/app/options/options.go'.
		// Please keep them in sync when doing update.
		CloudControllerManagerConfiguration: ccfg.CloudControllerManagerConfiguration{
			Generic: kubectrlmgrconfig.GenericControllerManagerConfiguration{
				Port:            ports.CloudControllerManagerPort,
				Address:         "0.0.0.0",
				MinResyncPeriod: metav1.Duration{Duration: 5 * time.Minute},
				ClientConnection: componentbaseconfig.ClientConnectionConfiguration{
					ContentType: "application/vnd.kubernetes.protobuf",
					QPS:         20.0,
					Burst:       30,
				},
				LeaderElection: componentbaseconfig.LeaderElectionConfiguration{
					LeaderElect:   false,
					LeaseDuration: metav1.Duration{Duration: 15 * time.Second},
					RenewDeadline: metav1.Duration{Duration: 10 * time.Second},
					RetryPeriod:   metav1.Duration{Duration: 2 * time.Second},
				},
				ControllerStartInterval: metav1.Duration{Duration: 0 * time.Second},
			},
			KubeCloudShared: kubectrlmgrconfig.KubeCloudSharedConfiguration{
				NodeMonitorPeriod:         metav1.Duration{Duration: 5 * time.Second},
				ClusterName:               "kubernetes",
				ConfigureCloudRoutes:      true,
				RouteReconciliationPeriod: metav1.Duration{Duration: 10 * time.Second},
			},
			ServiceController: serviceconfig.ServiceControllerConfiguration{
				ConcurrentServiceSyncs: 3,
			},
		},
		NodeStatusUpdateFrequency: metav1.Duration{Duration: 5 * time.Minute},
	}
	ccm.Generic.LeaderElection.LeaderElect = true
	return &ccm
}

func createRecorder(client *kubernetes.Clientset) record.EventRecorder {
	cast := record.NewBroadcaster()
	cast.StartLogging(klog.Infof)
	cast.StartRecordingToSink(
		&v1c.EventSinkImpl{
			Interface: v1c.New(client.CoreV1().RESTClient()).Events(""),
		},
	)
	return cast.NewRecorder(
		legacyscheme.Scheme,
		v1.EventSource{Component: "cloud-controller-manager"},
	)
}

func client(ccm *ServerCCM) (*kubernetes.Clientset, *kubernetes.Clientset, *rest.Config, error) {
	kubeconfig, err := clientcmd.BuildConfigFromFlags(ccm.Master, ccm.Kubeconfig)
	if err != nil {
		return nil, nil, kubeconfig, err
	}

	kubeconfig.QPS = ccm.Generic.ClientConnection.QPS
	kubeconfig.Burst = int(ccm.Generic.ClientConnection.Burst)
	kubeconfig.ContentConfig.ContentType = ccm.Generic.ClientConnection.ContentType

	client, err := kubernetes.NewForConfig(
		rest.AddUserAgent(
			kubeconfig,
			"cloud-controller-manager",
		),
	)
	if err != nil {
		klog.Fatalf("invalid API configuration: %v", err)
	}
	leader := kubernetes.NewForConfigOrDie(
		rest.AddUserAgent(
			kubeconfig,
			"leader-election",
		),
	)
	return client, leader, kubeconfig, nil
}

func (ccm *ServerCCM) initialization() error {
	if ccm.KubeCloudShared.CloudProvider.Name == "" {
		return fmt.Errorf("--cloud-provider cannot be empty")
	}

	alicloud.CloudConfigFile = ccm.KubeCloudShared.CloudProvider.CloudConfigFile
	cloud, err := cloudprovider.InitCloudProvider(
		ccm.KubeCloudShared.CloudProvider.Name,
		ccm.KubeCloudShared.CloudProvider.CloudConfigFile,
	)
	if err != nil {
		return fmt.Errorf("cloud provider could not be initialized: %v", err)
	}

	ccm.cloud = cloud
	if !cloud.HasClusterID() {
		if ccm.KubeCloudShared.AllowUntaggedCloud {
			klog.Warning("detected a cluster without a ClusterID.  A ClusterID will " +
				"be required in the future.  Please tag your cluster to avoid any future issues")
		} else {
			return fmt.Errorf("no ClusterID found.  A ClusterID is required for the " +
				"cloud provider to function properly.  This check can be bypassed by " +
				"setting the allow-untagged-cloud option")
		}
	}

	cfg, err := configz.New("componentconfig")
	if err != nil {
		return fmt.Errorf("unable to register configz: %v", err)
	}
	cfg.Set(ccm.CloudControllerManagerConfiguration)

	ccm.client, ccm.election, ccm.restConfig, err = client(ccm)
	if err != nil {
		return fmt.Errorf("create client error: %s", err.Error())
	}

	if err := setResourceLock(ccm); err != nil {
		return fmt.Errorf("set resource lock error: %s", err.Error())
	}

	ccm.recorder = createRecorder(ccm.client)
	return err
}

func (ccm *ServerCCM) Start() error {
	if err := ccm.initialization(); err != nil {
		return fmt.Errorf("verify ccm config: %s", err.Error())
	}

	// Start the external controller manager server
	go func() {
		mux := http.NewServeMux()
		healthz.InstallHandler(mux)
		if ccm.Generic.Debugging.EnableProfiling {
			mux.HandleFunc("/debug/pprof/", pprof.Index)
			mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
			mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
			mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
			if ccm.Generic.Debugging.EnableContentionProfiling {
				runtime.SetBlockProfileRate(1)
			}
		}
		configz.InstallHandler(mux)
		metric.RegisterPrometheus()
		mux.Handle("/metrics", promhttp.Handler())
		server := &http.Server{
			Addr:    net.JoinHostPort(ccm.Generic.Address, strconv.Itoa(int(ccm.Generic.Port))),
			Handler: mux,
		}
		klog.Fatal(server.ListenAndServe())
	}()
	return nil
}

func (ccm *ServerCCM) MainLoop(ctx context.Context) {
	var (
		builder controller.ControllerClientBuilder
	)
	builder = controller.SimpleControllerClientBuilder{
		ClientConfig: ccm.restConfig,
	}
	if ccm.KubeCloudShared.UseServiceAccountCredentials {
		builder = controller.SAControllerClientBuilder{
			ClientConfig:         rest.AnonymousClientConfig(ccm.restConfig),
			CoreClient:           ccm.client.CoreV1(),
			AuthenticationClient: ccm.client.AuthenticationV1(),
			Namespace:            "kube-system",
		}
	}
	panic(fmt.Sprintf("unreachable: %v", RunControllers(ccm, builder, nil)))
}

// Run runs the ExternalCMServer.  This should never exit.
func Run(ccm *ServerCCM) error {
	if err := ccm.Start(); err != nil {
		return err
	}

	route.Options = route.RoutesOptions{
		ClusterCIDR:               ccm.KubeCloudShared.ClusterCIDR,
		AllocateNodeCIDRs:         ccm.KubeCloudShared.AllocateNodeCIDRs,
		ConfigCloudRoutes:         ccm.KubeCloudShared.ConfigureCloudRoutes,
		RouteReconciliationPeriod: ccm.KubeCloudShared.RouteReconciliationPeriod,
		ControllerStartInterval:   ccm.Generic.ControllerStartInterval,
	}

	if !ccm.Generic.LeaderElection.LeaderElect {
		ccm.MainLoop(context.TODO())
	}

	// Identity used to distinguish between multiple cloud controller manager instances
	id, err := os.Hostname()
	if err != nil {
		return fmt.Errorf("hostname: %s", err.Error())
	}

	// Lock required for leader election
	rl, err := resourcelock.New(
		ccm.Generic.LeaderElection.ResourceLock,
		"kube-system",
		"ccm",
		ccm.election.CoreV1(),
		ccm.election.CoordinationV1(),
		resourcelock.ResourceLockConfig{
			Identity:      "ccm-" + id,
			EventRecorder: ccm.recorder,
		})
	if err != nil {
		return fmt.Errorf("creating etcd resource lock: %v", err)
	}

	// Try and become the leader and start cloud controller manager loops
	leaderelection.RunOrDie(
		context.Background(),
		leaderelection.LeaderElectionConfig{
			Lock: rl,
			Callbacks: leaderelection.LeaderCallbacks{
				OnStartedLeading: ccm.MainLoop,
				OnStoppedLeading: func() {
					klog.Fatalf("leaderelection lost")
				},
			},
			LeaseDuration: ccm.Generic.LeaderElection.LeaseDuration.Duration,
			RenewDeadline: ccm.Generic.LeaderElection.RenewDeadline.Duration,
			RetryPeriod:   ccm.Generic.LeaderElection.RetryPeriod.Duration,
		},
	)
	panic("unreachable")
}

// RunControllers starts the cloud specific controller loops.
func RunControllers(
	ccm *ServerCCM,
	clientBuilder controller.ControllerClientBuilder,
	stop <-chan struct{},
) error {

	if ccm.cloud != nil {
		// Initialize the cloud provider with a reference to the clientBuilder
		ccm.cloud.Initialize(clientBuilder, stop)
	}
	client := clientBuilder.ClientOrDie("shared-informers")

	ifactory := informers.NewSharedInformerFactory(client, resyncPeriod(ccm)())

	//if err := runControllerPV(ccm, clientBuilder, stop); err != nil {
	//	return fmt.Errorf("run pvcontroller: %s", err.Error())
	//}
	time.Sleep(wait.Jitter(ccm.Generic.ControllerStartInterval.Duration, ControllerStartJitter))

	if err := runControllerService(ccm, clientBuilder, ifactory, stop); err != nil {
		return fmt.Errorf("run service controller: %s", err.Error())
	}

	time.Sleep(wait.Jitter(ccm.Generic.ControllerStartInterval.Duration, ControllerStartJitter))

	// If apiserver is not running we should wait for some time and fail
	// only then. This is particularly important when we start apiserver
	// and controller manager at the same time.
	err := genericcontrollermanager.WaitForAPIServer(ccm.client, 10*time.Second)
	if err != nil {
		return fmt.Errorf("failed to get api versions from server: %v", err)
	}

	ifactory.Start(stop)
	klog.Infof("informer started")

	select {}
}

//func runControllerPV(
//	ccm *ServerCCM,
//	builder controller.ControllerClientBuilder,
//	stop <-chan struct{},
//) error {
//
//	con := ccontroller.NewPersistentVolumeLabelController(
//		builder.ClientOrDie("pvl-controller"),
//		ccm.cloud,
//	)
//	go con.Run(5, stop)
//	return nil
//}

func runControllerService(
	ccm *ServerCCM,
	builder controller.ControllerClientBuilder,
	informer informers.SharedInformerFactory,
	stop <-chan struct{},
) error {
	cloudslb, ok := ccm.cloud.LoadBalancer()
	if !ok {
		return fmt.Errorf("loadbalancer interface must be implemented")
	}

	scon, err := service.NewController(
		cloudslb,
		builder.ClientOrDie("cloud-controller-manager"),
		informer,
		ccm.KubeCloudShared.ClusterName,
	)
	if err != nil {
		return fmt.Errorf("failed to start service controller: %v", err)
	}
	go scon.Run(stop, int(ccm.ServiceController.ConcurrentServiceSyncs))
	return nil
}

func resyncPeriod(ccm *ServerCCM) func() time.Duration {
	return func() time.Duration {
		factor := rand.Float64() + 1
		return time.Duration(float64(ccm.Generic.MinResyncPeriod.Nanoseconds()) * factor)
	}
}

func setResourceLock(ccm *ServerCCM) error {
	// check kubernetes version to use lease or not
	version114, _ := version.ParseGeneric("v1.14.0")
	serverVersion, err := ccm.client.Discovery().ServerVersion()
	if err != nil {
		return fmt.Errorf("get kubernetes version error: %s", err.Error())
	}

	runningVersion, err := version.ParseGeneric(serverVersion.String())
	if err != nil {
		return fmt.Errorf("unexpected error parsing running Kubernetes version, %s", err.Error())
	}

	if runningVersion.AtLeast(version114) {
		ccm.CloudControllerManagerConfiguration.Generic.LeaderElection.ResourceLock = resourcelock.LeasesResourceLock
	} else {
		klog.V(5).Infof("kubernetes version: %s, resource lock use EndpointsResourceLock", runningVersion)
		ccm.CloudControllerManagerConfiguration.Generic.LeaderElection.ResourceLock = resourcelock.EndpointsResourceLock
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2017 All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package main

import (
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/hyperledger/fabric/bccsp/factory"
	"github.com/hyperledger/fabric/core/config/configtest"
	"github.com/hyperledger/fabric/internal/configtxgen/genesisconfig"
	"github.com/stretchr/testify/assert"
)

var tmpDir string

func TestMain(m *testing.M) {
	dir, err := ioutil.TempDir("", "configtxgen")
	if err != nil {
		panic("Error creating temp dir")
	}
	tmpDir = dir
	testResult := m.Run()
	os.RemoveAll(dir)

	os.Exit(testResult)
}

func TestInspectMissing(t *testing.T) {
	assert.Error(t, doInspectBlock("NonSenseBlockFileThatDoesn'tActuallyExist"), "Missing block")
}

func TestInspectBlock(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")

	config := genesisconfig.Load(genesisconfig.SampleInsecureSoloProfile, configtest.GetDevConfigDir())

	assert.NoError(t, doOutputBlock(config, "foo", blockDest), "Good block generation request")
	assert.NoError(t, doInspectBlock(blockDest), "Good block inspection request")
}

func TestMissingOrdererSection(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")

	config := genesisconfig.Load(genesisconfig.SampleInsecureSoloProfile, configtest.GetDevConfigDir())
	config.Orderer = nil

	assert.Error(t, doOutputBlock(config, "foo", blockDest), "Missing orderer section")
}

func TestMissingConsortiumSection(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")

	config := genesisconfig.Load(genesisconfig.SampleInsecureSoloProfile, configtest.GetDevConfigDir())
	config.Consortiums = nil

	assert.NoError(t, doOutputBlock(config, "foo", blockDest), "Missing consortiums section")
}

func TestMissingConsortiumValue(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())
	config.Consortium = ""

	assert.Error(t, doOutputChannelCreateTx(config, nil, "foo", configTxDest), "Missing Consortium value in Application Profile definition")
}

func TestMissingApplicationValue(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())
	config.Application = nil

	assert.Error(t, doOutputChannelCreateTx(config, nil, "foo", configTxDest), "Missing Application value in Application Profile definition")
}

func TestInspectMissingConfigTx(t *testing.T) {
	assert.Error(t, doInspectChannelCreateTx("ChannelCreateTxFileWhichDoesn'tReallyExist"), "Missing channel create tx file")
}

func TestInspectConfigTx(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())

	assert.NoError(t, doOutputChannelCreateTx(config, nil, "foo", configTxDest), "Good outputChannelCreateTx generation request")
	assert.NoError(t, doInspectChannelCreateTx(configTxDest), "Good configtx inspection request")
}

func TestGenerateAnchorPeersUpdate(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "anchorPeerUpdate")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())

	assert.NoError(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, genesisconfig.SampleOrgName), "Good anchorPeerUpdate request")
}

func TestBadAnchorPeersUpdates(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "anchorPeerUpdate")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())

	assert.Error(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, ""), "Bad anchorPeerUpdate request - asOrg empty")

	backupApplication := config.Application
	config.Application = nil
	assert.Error(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, genesisconfig.SampleOrgName), "Bad anchorPeerUpdate request")
	config.Application = backupApplication

	config.Application.Organizations[0] = &genesisconfig.Organization{Name: "FakeOrg", ID: "FakeOrg"}
	assert.Error(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, genesisconfig.SampleOrgName), "Bad anchorPeerUpdate request - fake org")
}

func TestConfigTxFlags(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")
	configTxDestAnchorPeers := filepath.Join(tmpDir, "configtxAnchorPeers")

	oldArgs := os.Args
	defer func() {
		os.Args = oldArgs
		flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	}()

	cleanup := configtest.SetDevFabricConfigPath(t)
	defer cleanup()
	devConfigDir := configtest.GetDevConfigDir()

	os.Args = []string{
		"cmd",
		"-channelID=testchannelid",
		"-outputCreateChannelTx=" + configTxDest,
		"-channelCreateTxBaseProfile=" + genesisconfig.SampleSingleMSPSoloProfile,
		"-profile=" + genesisconfig.SampleSingleMSPChannelProfile,
		"-configPath=" + devConfigDir,
		"-inspectChannelCreateTx=" + configTxDest,
		"-outputAnchorPeersUpdate=" + configTxDestAnchorPeers,
		"-asOrg=" + genesisconfig.SampleOrgName,
	}

	main()

	_, err := os.Stat(configTxDest)
	assert.NoError(t, err, "Configtx file is written successfully")
	_, err = os.Stat(configTxDestAnchorPeers)
	assert.NoError(t, err, "Configtx anchor peers file is written successfully")
}

func TestBlockFlags(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")
	oldArgs := os.Args
	defer func() {
		os.Args = oldArgs
		flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	}()
	os.Args = []string{
		"cmd",
		"-channelID=testchannelid",
		"-profile=" + genesisconfig.SampleSingleMSPSoloProfile,
		"-outputBlock=" + blockDest,
		"-inspectBlock=" + blockDest,
	}
	cleanup := configtest.SetDevFabricConfigPath(t)
	defer cleanup()

	main()

	_, err := os.Stat(blockDest)
	assert.NoError(t, err, "Block file is written successfully")
}

func TestPrintOrg(t *testing.T) {
	factory.InitFactories(nil)
	config := genesisconfig.LoadTopLevel(configtest.GetDevConfigDir())

	assert.NoError(t, doPrintOrg(config, genesisconfig.SampleOrgName), "Good org to print")

	err := doPrintOrg(config, genesisconfig.SampleOrgName+".wrong")
	assert.Error(t, err, "Bad org name")
	assert.Regexp(t, "organization [^ ]* not found", err.Error())

	config.Organizations[0] = &genesisconfig.Organization{Name: "FakeOrg", ID: "FakeOrg"}
	err = doPrintOrg(config, "FakeOrg")
	assert.Error(t, err, "Fake org")
	assert.Regexp(t, "bad org definition", err.Error())
}
<commit_msg>Improve configtxgen test coverage for unhappy paths (#1031)<commit_after>/*
Copyright IBM Corp. 2017 All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package main

import (
	"flag"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/hyperledger/fabric/bccsp/factory"
	"github.com/hyperledger/fabric/core/config/configtest"
	"github.com/hyperledger/fabric/internal/configtxgen/genesisconfig"
	"github.com/stretchr/testify/assert"
)

var tmpDir string

func TestMain(m *testing.M) {
	dir, err := ioutil.TempDir("", "configtxgen")
	if err != nil {
		panic("Error creating temp dir")
	}
	tmpDir = dir
	testResult := m.Run()
	os.RemoveAll(dir)

	os.Exit(testResult)
}

func TestInspectMissing(t *testing.T) {
	assert.Error(t, doInspectBlock("NonSenseBlockFileThatDoesn'tActuallyExist"), "Missing block")
}

func TestInspectBlock(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")

	config := genesisconfig.Load(genesisconfig.SampleInsecureSoloProfile, configtest.GetDevConfigDir())

	assert.NoError(t, doOutputBlock(config, "foo", blockDest), "Good block generation request")
	assert.NoError(t, doInspectBlock(blockDest), "Good block inspection request")
}

func TestInspectBlockErr(t *testing.T) {
	config := genesisconfig.Load(genesisconfig.SampleInsecureSoloProfile, configtest.GetDevConfigDir())

	assert.Error(t, doOutputBlock(config, "foo", ""), "Error writing genesis block:")
	assert.Error(t, doInspectBlock(""), "Could not read block")
}

func TestMissingOrdererSection(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")

	config := genesisconfig.Load(genesisconfig.SampleInsecureSoloProfile, configtest.GetDevConfigDir())
	config.Orderer = nil

	assert.Error(t, doOutputBlock(config, "foo", blockDest), "Missing orderer section")
}

func TestMissingConsortiumSection(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")

	config := genesisconfig.Load(genesisconfig.SampleInsecureSoloProfile, configtest.GetDevConfigDir())
	config.Consortiums = nil

	assert.NoError(t, doOutputBlock(config, "foo", blockDest), "Missing consortiums section")
}

func TestMissingConsortiumValue(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())
	config.Consortium = ""

	assert.Error(t, doOutputChannelCreateTx(config, nil, "foo", configTxDest), "Missing Consortium value in Application Profile definition")
}

func TestMissingApplicationValue(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())
	config.Application = nil

	assert.Error(t, doOutputChannelCreateTx(config, nil, "foo", configTxDest), "Missing Application value in Application Profile definition")
}

func TestInspectMissingConfigTx(t *testing.T) {
	assert.Error(t, doInspectChannelCreateTx("ChannelCreateTxFileWhichDoesn'tReallyExist"), "Missing channel create tx file")
}

func TestInspectConfigTx(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())

	assert.NoError(t, doOutputChannelCreateTx(config, nil, "foo", configTxDest), "Good outputChannelCreateTx generation request")
	assert.NoError(t, doInspectChannelCreateTx(configTxDest), "Good configtx inspection request")
}

func TestGenerateAnchorPeersUpdate(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "anchorPeerUpdate")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())

	assert.NoError(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, genesisconfig.SampleOrgName), "Good anchorPeerUpdate request")
}

func TestBadAnchorPeersUpdates(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "anchorPeerUpdate")

	config := genesisconfig.Load(genesisconfig.SampleSingleMSPChannelProfile, configtest.GetDevConfigDir())

	assert.Error(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, ""), "Bad anchorPeerUpdate request - asOrg empty")

	backupApplication := config.Application
	config.Application = nil
	assert.Error(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, genesisconfig.SampleOrgName), "Bad anchorPeerUpdate request")
	config.Application = backupApplication

	config.Application.Organizations[0] = &genesisconfig.Organization{Name: "FakeOrg", ID: "FakeOrg"}
	assert.Error(t, doOutputAnchorPeersUpdate(config, "foo", configTxDest, genesisconfig.SampleOrgName), "Bad anchorPeerUpdate request - fake org")
}

func TestConfigTxFlags(t *testing.T) {
	configTxDest := filepath.Join(tmpDir, "configtx")
	configTxDestAnchorPeers := filepath.Join(tmpDir, "configtxAnchorPeers")

	oldArgs := os.Args
	defer func() {
		os.Args = oldArgs
		flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	}()

	cleanup := configtest.SetDevFabricConfigPath(t)
	defer cleanup()
	devConfigDir := configtest.GetDevConfigDir()

	os.Args = []string{
		"cmd",
		"-channelID=testchannelid",
		"-outputCreateChannelTx=" + configTxDest,
		"-channelCreateTxBaseProfile=" + genesisconfig.SampleSingleMSPSoloProfile,
		"-profile=" + genesisconfig.SampleSingleMSPChannelProfile,
		"-configPath=" + devConfigDir,
		"-inspectChannelCreateTx=" + configTxDest,
		"-outputAnchorPeersUpdate=" + configTxDestAnchorPeers,
		"-asOrg=" + genesisconfig.SampleOrgName,
	}

	main()

	_, err := os.Stat(configTxDest)
	assert.NoError(t, err, "Configtx file is written successfully")
	_, err = os.Stat(configTxDestAnchorPeers)
	assert.NoError(t, err, "Configtx anchor peers file is written successfully")
}

func TestBlockFlags(t *testing.T) {
	blockDest := filepath.Join(tmpDir, "block")
	oldArgs := os.Args
	defer func() {
		os.Args = oldArgs
		flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	}()
	os.Args = []string{
		"cmd",
		"-channelID=testchannelid",
		"-profile=" + genesisconfig.SampleSingleMSPSoloProfile,
		"-outputBlock=" + blockDest,
		"-inspectBlock=" + blockDest,
	}
	cleanup := configtest.SetDevFabricConfigPath(t)
	defer cleanup()

	main()

	_, err := os.Stat(blockDest)
	assert.NoError(t, err, "Block file is written successfully")
}

func TestPrintOrg(t *testing.T) {
	factory.InitFactories(nil)
	config := genesisconfig.LoadTopLevel(configtest.GetDevConfigDir())

	assert.NoError(t, doPrintOrg(config, genesisconfig.SampleOrgName), "Good org to print")

	err := doPrintOrg(config, genesisconfig.SampleOrgName+".wrong")
	assert.Error(t, err, "Bad org name")
	assert.Regexp(t, "organization [^ ]* not found", err.Error())

	config.Organizations[0] = &genesisconfig.Organization{Name: "FakeOrg", ID: "FakeOrg"}
	err = doPrintOrg(config, "FakeOrg")
	assert.Error(t, err, "Fake org")
	assert.Regexp(t, "bad org definition", err.Error())
}
<|endoftext|>


<commit_before>// Copyright 2019 The Ebiten Authors
//
// 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 main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
)

const gomobileHash = "4a8be17bd2e3f4793a871839e766cbea7e999464"

func runCommand(command string, args []string, env []string) error {
	if buildX || buildN {
		for _, e := range env {
			fmt.Printf("%s ", e)
		}
		fmt.Print(command)
		for _, arg := range args {
			fmt.Printf(" %s", arg)
		}
		fmt.Println()
	}

	if buildN {
		return nil
	}

	cmd := exec.Command(command, args...)
	if len(env) > 0 {
		cmd.Env = append(os.Environ(), env...)
	}
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%s %v failed: %v\n%v", command, args, string(out), err)
	}
	return nil
}

func removeAll(path string) error {
	if buildX || buildN {
		fmt.Printf("rm -rf %s\n", path)
	}
	if buildN {
		return nil
	}
	return os.RemoveAll(path)
}

func runGo(args ...string) error {
	// TODO: Remove this after Ebiten drops the support of Go 1.15 and older.
	// GO111MODULE is on by default as of Go 1.16.
	env := []string{
		"GO111MODULE=on",
	}
	return runCommand("go", args, env)
}

// exe adds the .exe extension to the given filename.
// Without .exe, the executable won't be found by exec.LookPath on Windows (#1096).
func exe(filename string) string {
	if runtime.GOOS == "windows" {
		return filename + ".exe"
	}
	return filename
}

func prepareGomobileCommands() (string, error) {
	tmp, err := ioutil.TempDir("", "ebitenmobile-")
	if err != nil {
		return "", err
	}

	newpath := filepath.Join(tmp, "bin")
	if path := os.Getenv("PATH"); path != "" {
		newpath += string(filepath.ListSeparator) + path
	}
	if buildX || buildN {
		fmt.Printf("PATH=%s\n", newpath)
	}
	if !buildN {
		if err := os.Setenv("PATH", newpath); err != nil {
			return tmp, err
		}
	}

	pwd, err := os.Getwd()
	if err != nil {
		return tmp, err
	}

	// cd
	if buildX {
		fmt.Printf("cd %s\n", tmp)
	}
	if err := os.Chdir(tmp); err != nil {
		return tmp, err
	}
	defer func() {
		os.Chdir(pwd)
	}()

	const (
		modname   = "ebitenmobiletemporary"
		buildtags = "//go:build tools" +
			"\n// +build tools"
	)
	if err := runGo("mod", "init", modname); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile("tools.go", []byte(fmt.Sprintf(`%s

package %s

import (
	_ "golang.org/x/mobile/cmd/gobind"
	_ "golang.org/x/mobile/cmd/gomobile"
)
`, buildtags, modname)), 0644); err != nil {
		return tmp, err
	}

	// To record gomobile to go.sum for Go 1.16 and later, go-get gomobile instaed of golang.org/x/mobile (#1487).
	// This also records gobind as gomobile depends on gobind indirectly.
	// Using `...` doesn't work on Windows since mobile/internal/mobileinit cannot be compiled on Windows w/o Cgo (#1493).
	if err := runGo("get", "golang.org/x/mobile/cmd/gomobile@"+gomobileHash); err != nil {
		return tmp, err
	}
	if localgm := os.Getenv("EBITENMOBILE_GOMOBILE"); localgm != "" {
		if !filepath.IsAbs(localgm) {
			localgm = filepath.Join(pwd, localgm)
		}
		if err := runGo("mod", "edit", "-replace=golang.org/x/mobile="+localgm); err != nil {
			return tmp, err
		}
	}
	if err := runGo("mod", "tidy"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gomobile")), "golang.org/x/mobile/cmd/gomobile"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind-original")), "golang.org/x/mobile/cmd/gobind"); err != nil {
		return tmp, err
	}

	if err := os.Mkdir("src", 0755); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile(filepath.Join("src", "gobind.go"), gobindsrc, 0644); err != nil {
		return tmp, err
	}

	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind")), "-tags", "ebitenmobilegobind", filepath.Join("src", "gobind.go")); err != nil {
		return tmp, err
	}

	if err := runCommand("gomobile", []string{"init"}, nil); err != nil {
		return tmp, err
	}

	return tmp, nil
}
<commit_msg>cmd/ebitenmobile: adjust gomobile version<commit_after>// Copyright 2019 The Ebiten Authors
//
// 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 main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
)

const gomobileHash = "447654d348e36cdac65810aa73b4fc8e7b3ceef1"

func runCommand(command string, args []string, env []string) error {
	if buildX || buildN {
		for _, e := range env {
			fmt.Printf("%s ", e)
		}
		fmt.Print(command)
		for _, arg := range args {
			fmt.Printf(" %s", arg)
		}
		fmt.Println()
	}

	if buildN {
		return nil
	}

	cmd := exec.Command(command, args...)
	if len(env) > 0 {
		cmd.Env = append(os.Environ(), env...)
	}
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%s %v failed: %v\n%v", command, args, string(out), err)
	}
	return nil
}

func removeAll(path string) error {
	if buildX || buildN {
		fmt.Printf("rm -rf %s\n", path)
	}
	if buildN {
		return nil
	}
	return os.RemoveAll(path)
}

func runGo(args ...string) error {
	// TODO: Remove this after Ebiten drops the support of Go 1.15 and older.
	// GO111MODULE is on by default as of Go 1.16.
	env := []string{
		"GO111MODULE=on",
	}
	return runCommand("go", args, env)
}

// exe adds the .exe extension to the given filename.
// Without .exe, the executable won't be found by exec.LookPath on Windows (#1096).
func exe(filename string) string {
	if runtime.GOOS == "windows" {
		return filename + ".exe"
	}
	return filename
}

func prepareGomobileCommands() (string, error) {
	tmp, err := ioutil.TempDir("", "ebitenmobile-")
	if err != nil {
		return "", err
	}

	newpath := filepath.Join(tmp, "bin")
	if path := os.Getenv("PATH"); path != "" {
		newpath += string(filepath.ListSeparator) + path
	}
	if buildX || buildN {
		fmt.Printf("PATH=%s\n", newpath)
	}
	if !buildN {
		if err := os.Setenv("PATH", newpath); err != nil {
			return tmp, err
		}
	}

	pwd, err := os.Getwd()
	if err != nil {
		return tmp, err
	}

	// cd
	if buildX {
		fmt.Printf("cd %s\n", tmp)
	}
	if err := os.Chdir(tmp); err != nil {
		return tmp, err
	}
	defer func() {
		os.Chdir(pwd)
	}()

	const (
		modname   = "ebitenmobiletemporary"
		buildtags = "//go:build tools" +
			"\n// +build tools"
	)
	if err := runGo("mod", "init", modname); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile("tools.go", []byte(fmt.Sprintf(`%s

package %s

import (
	_ "golang.org/x/mobile/cmd/gobind"
	_ "golang.org/x/mobile/cmd/gomobile"
)
`, buildtags, modname)), 0644); err != nil {
		return tmp, err
	}

	// To record gomobile to go.sum for Go 1.16 and later, go-get gomobile instaed of golang.org/x/mobile (#1487).
	// This also records gobind as gomobile depends on gobind indirectly.
	// Using `...` doesn't work on Windows since mobile/internal/mobileinit cannot be compiled on Windows w/o Cgo (#1493).
	if err := runGo("get", "golang.org/x/mobile/cmd/gomobile@"+gomobileHash); err != nil {
		return tmp, err
	}
	if localgm := os.Getenv("EBITENMOBILE_GOMOBILE"); localgm != "" {
		if !filepath.IsAbs(localgm) {
			localgm = filepath.Join(pwd, localgm)
		}
		if err := runGo("mod", "edit", "-replace=golang.org/x/mobile="+localgm); err != nil {
			return tmp, err
		}
	}
	if err := runGo("mod", "tidy"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gomobile")), "golang.org/x/mobile/cmd/gomobile"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind-original")), "golang.org/x/mobile/cmd/gobind"); err != nil {
		return tmp, err
	}

	if err := os.Mkdir("src", 0755); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile(filepath.Join("src", "gobind.go"), gobindsrc, 0644); err != nil {
		return tmp, err
	}

	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind")), "-tags", "ebitenmobilegobind", filepath.Join("src", "gobind.go")); err != nil {
		return tmp, err
	}

	if err := runCommand("gomobile", []string{"init"}, nil); err != nil {
		return tmp, err
	}

	return tmp, nil
}
<|endoftext|>


<commit_before>package api

import (
	"time"
)

type Config struct {
	AppId  string
	Secret string

	Timeout   time.Duration
	KeepAlive time.Duration

	Sub struct {
		Scheme   string // https or http
		Endpoint string // host:port
	}

	Pub struct {
		Scheme   string // http or https
		Endpoint string // host:port
	}

	Admin struct {
		Scheme   string
		Endpoint string
	}

	Debug bool
}

func DefaultConfig(appid, secret string) *Config {
	cf := &Config{
		AppId:     appid,
		Secret:    secret,
		Timeout:   time.Second * 120, // FIXME
		KeepAlive: time.Minute,
		Debug:     false,
	}
	cf.Sub.Scheme = "http"
	cf.Pub.Scheme = "http"
	cf.Admin.Scheme = "http"
	return cf
}
<commit_msg>pubsub api default timeout=1m<commit_after>package api

import (
	"time"
)

type Config struct {
	AppId  string
	Secret string

	Timeout   time.Duration
	KeepAlive time.Duration

	Sub struct {
		Scheme   string // https or http
		Endpoint string // host:port
	}

	Pub struct {
		Scheme   string // http or https
		Endpoint string // host:port
	}

	Admin struct {
		Scheme   string
		Endpoint string
	}

	Debug bool
}

func DefaultConfig(appid, secret string) *Config {
	cf := &Config{
		AppId:     appid,
		Secret:    secret,
		Timeout:   time.Minute,
		KeepAlive: time.Minute,
		Debug:     false,
	}
	cf.Sub.Scheme = "http"
	cf.Pub.Scheme = "http"
	cf.Admin.Scheme = "http"
	return cf
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/gocql/gocql"
	"github.com/raintank/metrictank/cluster"
	"github.com/raintank/worldping-api/pkg/log"
	"gopkg.in/raintank/schema.v1"
)

const table_schema = `CREATE TABLE IF NOT EXISTS %s.metric_idx (
    id text,
    partition int,
    name text,
    metric text,
    interval int,
    unit text,
    mtype text,
    tags set<text>,
    lastupdate int,
    PRIMARY KEY (id, partition)
) WITH compaction = {'class': 'SizeTieredCompactionStrategy'}
    AND compression = {'sstable_compression': 'org.apache.cassandra.io.compress.LZ4Compressor'}`
const metric_idx_index = `CREATE INDEX IF NOT EXISTS ON %s.metric_idx(partition)`

var (
	dryRun        = flag.Bool("dry-run", true, "run in dry-run mode. No changes will be made.")
	logLevel      = flag.Int("log-level", 2, "log level. 0=TRACE|1=DEBUG|2=INFO|3=WARN|4=ERROR|5=CRITICAL|6=FATAL")
	cassAddr      = flag.String("cass-addr", "localhost", "Address of cassandra host.")
	keyspace      = flag.String("keyspace", "raintank", "Cassandra keyspace to use.")
	partitionBy   = flag.String("partition-by", "byOrg", "method used for paritioning metrics. (byOrg|bySeries)")
	numPartitions = flag.Int("num-partitions", 1, "number of partitions in cluster")
)

func main() {
	flag.Parse()
	log.NewLogger(0, "console", fmt.Sprintf(`{"level": %d, "formatting":false}`, *logLevel))

	defsChan := make(chan *schema.MetricDefinition, 100)
	var wg sync.WaitGroup

	cluster := gocql.NewCluster(*cassAddr)
	cluster.Consistency = gocql.ParseConsistency("one")
	cluster.Timeout = time.Second
	cluster.NumConns = 2
	cluster.ProtoVersion = 4
	cluster.Keyspace = *keyspace
	session, err := cluster.CreateSession()
	if err != nil {
		log.Fatal(4, "failed to create cql session. %s", err)
	}

	// ensure the new table exists.
	err = session.Query(fmt.Sprintf(table_schema, *keyspace)).Exec()
	if err != nil {
		log.Fatal(4, "cassandra-idx failed to initialize cassandra table. %s", err)
	}
	err = session.Query(fmt.Sprintf(metric_idx_index, *keyspace)).Exec()
	if err != nil {
		log.Fatal(4, "cassandra-idx failed to initialize cassandra index. %s", err)
	}

	wg.Add(1)
	go writeDefs(session, defsChan, wg)
	wg.Add(1)
	go getDefs(session, defsChan, wg)

	wg.Wait()

}

func writeDefs(session *gocql.Session, defsChan chan *schema.MetricDefinition, wg sync.WaitGroup) {
	log.Info("starting write thread")
	defer wg.Done()
	for def := range defsChan {
		qry := `INSERT INTO metric_idx (id, partition, name, metric, interval, unit, mtype, tags, lastupdate) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`
		if *dryRun {
			fmt.Printf(
				"INSERT INTO metric_idx (id, partition, name, metric, interval, unit, mtype, tags, lastupdate) VALUES ('%s', '%d','%s', '%s','%d', '%s','%s', '%v', '%d')\n",
				def.Id,
				def.Partition,
				def.Name,
				def.Metric,
				def.Interval,
				def.Unit,
				def.Mtype,
				def.Tags,
				def.LastUpdate)
			continue
		}
		success := false
		attempts := 0
		for !success {
			if err := session.Query(
				qry,
				def.Id,
				def.Partition,
				def.Name,
				def.Metric,
				def.Interval,
				def.Unit,
				def.Mtype,
				def.Tags,
				def.LastUpdate).Exec(); err != nil {

				if (attempts % 20) == 0 {
					log.Warn("cassandra-idx Failed to write def to cassandra. it will be retried. %s", err)
				}
				sleepTime := 100 * attempts
				if sleepTime > 2000 {
					sleepTime = 2000
				}
				time.Sleep(time.Duration(sleepTime) * time.Millisecond)
				attempts++
			} else {
				success = true
				log.Debug("cassandra-idx metricDef saved to cassandra. %s", def.Id)
			}
		}
	}
}

func getDefs(session *gocql.Session, defsChan chan *schema.MetricDefinition, wg sync.WaitGroup) {
	log.Info("starting read thread")
	defer wg.Done()
	defer close(defsChan)
	partitioner, err := cluster.NewKafkaPartitioner(*partitionBy)
	if err != nil {
		log.Fatal(4, "failed to initialize partitioner. %s", err)
	}
	iter := session.Query("SELECT def from metric_def_idx").Iter()

	var data []byte
	for iter.Scan(&data) {
		mdef := schema.MetricDefinition{}
		_, err := mdef.UnmarshalMsg(data)
		if err != nil {
			log.Error(3, "cassandra-idx Bad definition in index. %s - %s", data, err)
			continue
		}
		log.Debug("retrieved %s from old index.", mdef.Id)
		if *numPartitions == 1 {
			def.Partition = 0
		} else {
			p, err := partitioner.Partition(&mdef, *numPartitions)
			if err != nil {
				log.Error(3, "failed to get partition id of metric. %s", err)
				mdef.Partition = 0
			} else {
				mdef.Partition = p
			}
		}
		defsChan <- &mdef
	}
}
<commit_msg>use updated metric_idx schema<commit_after>package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/gocql/gocql"
	"github.com/raintank/metrictank/cluster"
	"github.com/raintank/worldping-api/pkg/log"
	"gopkg.in/raintank/schema.v1"
)

const table_schema = `CREATE TABLE IF NOT EXISTS %s.metric_idx (
    id text,
    orgId int,
    partition int,
    name text,
    metric text,
    interval int,
    unit text,
    mtype text,
    tags set<text>,
    lastupdate int,
    PRIMARY KEY (id, partition)
) WITH compaction = {'class': 'SizeTieredCompactionStrategy'}
    AND compression = {'sstable_compression': 'org.apache.cassandra.io.compress.LZ4Compressor'}`
const metric_idx_index = `CREATE INDEX IF NOT EXISTS ON %s.metric_idx(partition)`

var (
	dryRun        = flag.Bool("dry-run", true, "run in dry-run mode. No changes will be made.")
	logLevel      = flag.Int("log-level", 2, "log level. 0=TRACE|1=DEBUG|2=INFO|3=WARN|4=ERROR|5=CRITICAL|6=FATAL")
	cassAddr      = flag.String("cass-addr", "localhost", "Address of cassandra host.")
	keyspace      = flag.String("keyspace", "raintank", "Cassandra keyspace to use.")
	partitionBy   = flag.String("partition-by", "byOrg", "method used for paritioning metrics. (byOrg|bySeries)")
	numPartitions = flag.Int("num-partitions", 1, "number of partitions in cluster")
)

func main() {
	flag.Parse()
	log.NewLogger(0, "console", fmt.Sprintf(`{"level": %d, "formatting":false}`, *logLevel))

	defsChan := make(chan *schema.MetricDefinition, 100)
	var wg sync.WaitGroup

	cluster := gocql.NewCluster(*cassAddr)
	cluster.Consistency = gocql.ParseConsistency("one")
	cluster.Timeout = time.Second
	cluster.NumConns = 2
	cluster.ProtoVersion = 4
	cluster.Keyspace = *keyspace
	session, err := cluster.CreateSession()
	if err != nil {
		log.Fatal(4, "failed to create cql session. %s", err)
	}

	// ensure the new table exists.
	err = session.Query(fmt.Sprintf(table_schema, *keyspace)).Exec()
	if err != nil {
		log.Fatal(4, "cassandra-idx failed to initialize cassandra table. %s", err)
	}
	err = session.Query(fmt.Sprintf(metric_idx_index, *keyspace)).Exec()
	if err != nil {
		log.Fatal(4, "cassandra-idx failed to initialize cassandra index. %s", err)
	}

	wg.Add(1)
	go writeDefs(session, defsChan, wg)
	wg.Add(1)
	go getDefs(session, defsChan, wg)

	wg.Wait()

}

func writeDefs(session *gocql.Session, defsChan chan *schema.MetricDefinition, wg sync.WaitGroup) {
	log.Info("starting write thread")
	defer wg.Done()
	for def := range defsChan {
		qry := `INSERT INTO metric_idx (id, orgid, partition, name, metric, interval, unit, mtype, tags, lastupdate) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`
		if *dryRun {
			fmt.Printf(
				"INSERT INTO metric_idx (id, orgid, partition, name, metric, interval, unit, mtype, tags, lastupdate) VALUES ('%s', '%d', '%d','%s', '%s','%d', '%s','%s', '%v', '%d')\n",
				def.Id,
				def.OrgId,
				def.Partition,
				def.Name,
				def.Metric,
				def.Interval,
				def.Unit,
				def.Mtype,
				def.Tags,
				def.LastUpdate)
			continue
		}
		success := false
		attempts := 0
		for !success {
			if err := session.Query(
				qry,
				def.Id,
				def.OrgId,
				def.Partition,
				def.Name,
				def.Metric,
				def.Interval,
				def.Unit,
				def.Mtype,
				def.Tags,
				def.LastUpdate).Exec(); err != nil {

				if (attempts % 20) == 0 {
					log.Warn("cassandra-idx Failed to write def to cassandra. it will be retried. %s", err)
				}
				sleepTime := 100 * attempts
				if sleepTime > 2000 {
					sleepTime = 2000
				}
				time.Sleep(time.Duration(sleepTime) * time.Millisecond)
				attempts++
			} else {
				success = true
				log.Debug("cassandra-idx metricDef saved to cassandra. %s", def.Id)
			}
		}
	}
}

func getDefs(session *gocql.Session, defsChan chan *schema.MetricDefinition, wg sync.WaitGroup) {
	log.Info("starting read thread")
	defer wg.Done()
	defer close(defsChan)
	partitioner, err := cluster.NewKafkaPartitioner(*partitionBy)
	if err != nil {
		log.Fatal(4, "failed to initialize partitioner. %s", err)
	}
	iter := session.Query("SELECT def from metric_def_idx").Iter()

	var data []byte
	for iter.Scan(&data) {
		mdef := schema.MetricDefinition{}
		_, err := mdef.UnmarshalMsg(data)
		if err != nil {
			log.Error(3, "cassandra-idx Bad definition in index. %s - %s", data, err)
			continue
		}
		log.Debug("retrieved %s from old index.", mdef.Id)
		if *numPartitions == 1 {
			def.Partition = 0
		} else {
			p, err := partitioner.Partition(&mdef, *numPartitions)
			if err != nil {
				log.Error(3, "failed to get partition id of metric. %s", err)
				mdef.Partition = 0
			} else {
				mdef.Partition = p
			}
		}
		defsChan <- &mdef
	}
}
<|endoftext|>