<commit_before>package state

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"

	"github.com/coreos/etcd/raft"
	"github.com/docker/swarm-v2/api"
	"github.com/stretchr/testify/assert"
)

const testTick = 50 * time.Millisecond

var (
	raftLogger = &raft.DefaultLogger{Logger: log.New(ioutil.Discard, "", 0)}
)

func init() {
	grpclog.SetLogger(log.New(ioutil.Discard, "", log.LstdFlags))
}

func pollNodeFunc(nodes map[uint64]*Node, f func(*Node) bool) error {
	var tickers []*time.Ticker
	var wg sync.WaitGroup
	for _, n := range nodes {
		wg.Add(1)
		if f(n) {
			wg.Done()
			continue
		}
		tick := time.NewTicker(testTick)
		tickers = append(tickers, tick)
		go func(n *Node) {
			for range tick.C {
				if f(n) {
					break
				}
			}
			wg.Done()
		}(n)
	}
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()
	select {
	case <-done:
		return nil
	case <-time.After(32 * time.Second):
		for _, tick := range tickers {
			tick.Stop()
		}
		return fmt.Errorf("polling failed")
	}
}

func getTermAndLeader(n *Node) (uint64, uint64) {
	status := n.Status()
	return status.Lead, status.Term
}

// waitForCluster waits until leader will be one of specified nodes
func waitForCluster(t *testing.T, nodes map[uint64]*Node) {
	assert.NoError(t, pollNodeFunc(nodes, func(n *Node) bool {
		l, t := getTermAndLeader(n)
		if _, ok := nodes[l]; !ok {
			return false
		}
		for _, node := range nodes {
			rl, rt := getTermAndLeader(node)
			if rl != l || rt != t {
				return false
			}
		}
		return true
	}))
}

// waitForPeerNumber waits until peers in cluster converge to specified number
func waitForPeerNumber(t *testing.T, nodes map[uint64]*Node, count int) {
	assert.NoError(t, pollNodeFunc(nodes, func(n *Node) bool {
		return len(n.Cluster.Peers()) == count
	}))
}

func newInitNode(t *testing.T, id uint64) *Node {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	assert.NoError(t, err, "can't bind to raft service port")
	s := grpc.NewServer()

	cfg := DefaultNodeConfig()
	cfg.Logger = raftLogger

	stateDir, err := ioutil.TempDir("", "test-raft")
	assert.NoError(t, err, "can't create temporary state directory")

	newNodeOpts := NewNodeOptions{
		ID:           id,
		Addr:         l.Addr().String(),
		Config:       cfg,
		StateDir:     stateDir,
		TickInterval: testTick,
	}
	n, err := NewNode(context.Background(), newNodeOpts)
	assert.NoError(t, err, "can't create raft node")
	n.Listener = l
	n.Server = s

	err = n.Campaign(n.Ctx)
	assert.NoError(t, err, "can't campaign to be the leader")
	n.Start()

	Register(s, n)

	done := make(chan error)
	go func() {
		done <- s.Serve(l)
	}()
	go func() {
		// After stopping, we should receive an error from Serve
		assert.Error(t, <-done)
	}()

	return n
}

func newJoinNode(t *testing.T, id uint64, join string) *Node {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	assert.NoError(t, err, "can't bind to raft service port")
	s := grpc.NewServer()

	cfg := DefaultNodeConfig()
	cfg.Logger = raftLogger

	stateDir, err := ioutil.TempDir("", "test-raft")
	assert.NoError(t, err, "can't create temporary state directory")

	newNodeOpts := NewNodeOptions{
		ID:           id,
		Addr:         l.Addr().String(),
		Config:       cfg,
		StateDir:     stateDir,
		TickInterval: testTick,
	}

	n, err := NewNode(context.Background(), newNodeOpts)
	assert.NoError(t, err, "can't create raft node")
	n.Listener = l
	n.Server = s

	n.Start()

	c, err := GetRaftClient(join, 500*time.Millisecond)
	assert.NoError(t, err, "can't initiate connection with existing raft")

	resp, err := c.Join(n.Ctx, &api.JoinRequest{
		Node: &api.RaftNode{ID: id, Addr: l.Addr().String()},
	})
	assert.NoError(t, err, "can't join existing Raft")

	err = n.RegisterNodes(resp.Members)
	assert.NoError(t, err, "can't add nodes to the local cluster list")

	Register(s, n)

	done := make(chan error)
	go func() {
		done <- s.Serve(l)
	}()
	go func() {
		// After stopping, we should receive an error from Serve
		assert.Error(t, <-done)
	}()

	return n
}

func newRaftCluster(t *testing.T) map[uint64]*Node {
	nodes := make(map[uint64]*Node)
	nodes[1] = newInitNode(t, 1)
	nodes[2] = newJoinNode(t, 2, nodes[1].Listener.Addr().String())
	nodes[3] = newJoinNode(t, 3, nodes[1].Listener.Addr().String())
	waitForCluster(t, nodes)
	return nodes
}

func addRaftNode(t *testing.T, nodes map[uint64]*Node) {
	n := uint64(len(nodes) + 1)
	nodes[n] = newJoinNode(t, uint64(n), nodes[1].Listener.Addr().String())
	waitForCluster(t, nodes)
}

func teardownCluster(t *testing.T, nodes map[uint64]*Node) {
	for _, node := range nodes {
		shutdownNode(node)
	}
}

func removeNode(nodes map[string]*Node, node string) {
	shutdownNode(nodes[node])
	delete(nodes, node)
}

func shutdownNode(node *Node) {
	node.Server.Stop()
	node.Shutdown()
	os.RemoveAll(node.stateDir)
}

func TestRaftBootstrap(t *testing.T) {
	t.Parallel()<commit_msg>raft: wait for cluster readiness on each added node<commit_after>package state

import (
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"testing"
	"time"

	"golang.org/x/net/context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"

	"github.com/coreos/etcd/raft"
	"github.com/docker/swarm-v2/api"
	"github.com/stretchr/testify/assert"
)

const testTick = 1 * time.Second

var (
	raftLogger = &raft.DefaultLogger{Logger: log.New(ioutil.Discard, "", 0)}
)

func init() {
	grpclog.SetLogger(log.New(ioutil.Discard, "", log.LstdFlags))
}

func pollFunc(f func() bool) error {
	if f() {
		return nil
	}
	tick := time.NewTicker(100 * time.Millisecond)
	defer tick.Stop()
	done := make(chan struct{})
	go func() {
		for range tick.C {
			if f() {
				break
			}
		}
		close(done)
	}()
	select {
	case <-done:
		return nil
	case <-time.After(32 * time.Second):
		return fmt.Errorf("polling failed")
	}
}

type leaderStatus struct {
	leader uint64
	term   uint64
}

func getTermAndLeader(n *Node) *leaderStatus {
	status := n.Status()
	return &leaderStatus{leader: status.Lead, term: status.Term}
}

// waitForCluster waits until leader will be one of specified nodes
func waitForCluster(t *testing.T, nodes map[uint64]*Node) {
	err := pollFunc(func() bool {
		var prev *leaderStatus
		for _, n := range nodes {
			if prev == nil {
				prev = getTermAndLeader(n)
				if _, ok := nodes[prev.leader]; !ok {
					return false
				}
				continue
			}
			cur := getTermAndLeader(n)
			if _, ok := nodes[cur.leader]; !ok {
				return false
			}
			if cur.leader != prev.leader || cur.term != prev.term {
				return false
			}
		}
		return true
	})
	assert.NoError(t, err)
}

// waitForPeerNumber waits until peers in cluster converge to specified number
func waitForPeerNumber(t *testing.T, nodes map[uint64]*Node, count int) {
	assert.NoError(t, pollFunc(func() bool {
		for _, n := range nodes {
			if len(n.Cluster.Peers()) != count {
				return false
			}
		}
		return true
	}))
}

func newInitNode(t *testing.T, id uint64) *Node {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	assert.NoError(t, err, "can't bind to raft service port")
	s := grpc.NewServer()

	cfg := DefaultNodeConfig()
	cfg.Logger = raftLogger

	stateDir, err := ioutil.TempDir("", "test-raft")
	assert.NoError(t, err, "can't create temporary state directory")

	newNodeOpts := NewNodeOptions{
		ID:           id,
		Addr:         l.Addr().String(),
		Config:       cfg,
		StateDir:     stateDir,
		TickInterval: testTick,
	}
	n, err := NewNode(context.Background(), newNodeOpts)
	assert.NoError(t, err, "can't create raft node")
	n.Listener = l
	n.Server = s

	err = n.Campaign(n.Ctx)
	assert.NoError(t, err, "can't campaign to be the leader")
	n.Start()

	Register(s, n)

	done := make(chan error)
	go func() {
		done <- s.Serve(l)
	}()
	go func() {
		// After stopping, we should receive an error from Serve
		assert.Error(t, <-done)
	}()

	return n
}

func newJoinNode(t *testing.T, id uint64, join string) *Node {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	assert.NoError(t, err, "can't bind to raft service port")
	s := grpc.NewServer()

	cfg := DefaultNodeConfig()
	cfg.Logger = raftLogger

	stateDir, err := ioutil.TempDir("", "test-raft")
	assert.NoError(t, err, "can't create temporary state directory")

	newNodeOpts := NewNodeOptions{
		ID:           id,
		Addr:         l.Addr().String(),
		Config:       cfg,
		StateDir:     stateDir,
		TickInterval: testTick,
	}

	n, err := NewNode(context.Background(), newNodeOpts)
	assert.NoError(t, err, "can't create raft node")
	n.Listener = l
	n.Server = s

	n.Start()

	c, err := GetRaftClient(join, 500*time.Millisecond)
	assert.NoError(t, err, "can't initiate connection with existing raft")

	resp, err := c.Join(n.Ctx, &api.JoinRequest{
		Node: &api.RaftNode{ID: id, Addr: l.Addr().String()},
	})
	assert.NoError(t, err, "can't join existing Raft")

	err = n.RegisterNodes(resp.Members)
	assert.NoError(t, err, "can't add nodes to the local cluster list")

	Register(s, n)

	done := make(chan error)
	go func() {
		done <- s.Serve(l)
	}()
	go func() {
		// After stopping, we should receive an error from Serve
		assert.Error(t, <-done)
	}()
	return n
}

func newRaftCluster(t *testing.T) map[uint64]*Node {
	nodes := make(map[uint64]*Node)
	nodes[1] = newInitNode(t, 1)
	addRaftNode(t, nodes)
	addRaftNode(t, nodes)
	return nodes
}

func addRaftNode(t *testing.T, nodes map[uint64]*Node) {
	n := uint64(len(nodes) + 1)
	nodes[n] = newJoinNode(t, uint64(n), nodes[1].Listener.Addr().String())
	waitForCluster(t, nodes)
}

func teardownCluster(t *testing.T, nodes map[uint64]*Node) {
	for _, node := range nodes {
		shutdownNode(node)
	}
}

func removeNode(nodes map[string]*Node, node string) {
	shutdownNode(nodes[node])
	delete(nodes, node)
}

func shutdownNode(node *Node) {
	node.Server.Stop()
	node.Shutdown()
	os.RemoveAll(node.stateDir)
}

func TestRaftBootstrap(t *testing.T) {
	t.Parallel()
{'fraction_non_alphanumeric': 0.15822976659593815, 'fraction_numerical': 0.011417601293321209, 'mean_word_length': 3.702137767220903, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 4, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1750794', 'n_tokens_mistral': 4066, 'n_tokens_neox': 3754, 'n_words': 915}