github.com/macb/etcd@v0.3.1-0.20140227003422-a60481c6b1a0/tests/server_utils.go (about)

     1  package tests
     2  
     3  import (
     4  	"io/ioutil"
     5  	"net/http"
     6  	"os"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/coreos/etcd/third_party/github.com/coreos/raft"
    11  
    12  	"github.com/coreos/etcd/metrics"
    13  	"github.com/coreos/etcd/server"
    14  	"github.com/coreos/etcd/store"
    15  )
    16  
    17  const (
    18  	testName		= "ETCDTEST"
    19  	testClientURL		= "localhost:4401"
    20  	testRaftURL		= "localhost:7701"
    21  	testSnapshotCount	= 10000
    22  	testHeartbeatInterval	= time.Duration(50) * time.Millisecond
    23  	testElectionTimeout	= time.Duration(200) * time.Millisecond
    24  )
    25  
    26  // Starts a server in a temporary directory.
    27  func RunServer(f func(*server.Server)) {
    28  	path, _ := ioutil.TempDir("", "etcd-")
    29  	defer os.RemoveAll(path)
    30  
    31  	store := store.New()
    32  	registry := server.NewRegistry(store)
    33  
    34  	serverStats := server.NewRaftServerStats(testName)
    35  	followersStats := server.NewRaftFollowersStats(testName)
    36  
    37  	psConfig := server.PeerServerConfig{
    38  		Name:		testName,
    39  		URL:		"http://" + testRaftURL,
    40  		Scheme:		"http",
    41  		SnapshotCount:	testSnapshotCount,
    42  		MaxClusterSize:	9,
    43  	}
    44  
    45  	mb := metrics.NewBucket("")
    46  
    47  	ps := server.NewPeerServer(psConfig, registry, store, &mb, followersStats, serverStats)
    48  	psListener, err := server.NewListener(testRaftURL)
    49  	if err != nil {
    50  		panic(err)
    51  	}
    52  
    53  	// Create Raft transporter and server
    54  	dialTimeout := (3 * testHeartbeatInterval) + testElectionTimeout
    55  	responseHeaderTimeout := (3 * testHeartbeatInterval) + testElectionTimeout
    56  	raftTransporter := server.NewTransporter(followersStats, serverStats, registry, testHeartbeatInterval, dialTimeout, responseHeaderTimeout)
    57  	raftServer, err := raft.NewServer(testName, path, raftTransporter, store, ps, "")
    58  	if err != nil {
    59  		panic(err)
    60  	}
    61  	raftServer.SetElectionTimeout(testElectionTimeout)
    62  	raftServer.SetHeartbeatInterval(testHeartbeatInterval)
    63  	ps.SetRaftServer(raftServer)
    64  
    65  	s := server.New(testName, "http://"+testClientURL, ps, registry, store, nil)
    66  	sListener, err := server.NewListener(testClientURL)
    67  	if err != nil {
    68  		panic(err)
    69  	}
    70  
    71  	ps.SetServer(s)
    72  
    73  	w := &sync.WaitGroup{}
    74  
    75  	// Start up peer server.
    76  	c := make(chan bool)
    77  	go func() {
    78  		c <- true
    79  		ps.Start(false, "", []string{})
    80  		h := waitHandler{w, ps.HTTPHandler()}
    81  		http.Serve(psListener, &h)
    82  	}()
    83  	<-c
    84  
    85  	// Start up etcd server.
    86  	go func() {
    87  		c <- true
    88  		h := waitHandler{w, s.HTTPHandler()}
    89  		http.Serve(sListener, &h)
    90  	}()
    91  	<-c
    92  
    93  	// Wait to make sure servers have started.
    94  	time.Sleep(50 * time.Millisecond)
    95  
    96  	// Execute the function passed in.
    97  	f(s)
    98  
    99  	// Clean up servers.
   100  	ps.Stop()
   101  	psListener.Close()
   102  	sListener.Close()
   103  	w.Wait()
   104  }
   105  
   106  type waitHandler struct {
   107          wg *sync.WaitGroup
   108          handler http.Handler
   109  }
   110  
   111  func (h *waitHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){
   112          h.wg.Add(1)
   113          defer h.wg.Done()
   114          h.handler.ServeHTTP(w, r)
   115  
   116          //important to flush before decrementing the wait group.
   117          //we won't get a chance to once main() ends.
   118          w.(http.Flusher).Flush()
   119  }