github.com/ghodss/etcd@v0.3.1-0.20140417172404-cc329bfa55cb/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/goraft/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  	}
    43  
    44  	mb := metrics.NewBucket("")
    45  
    46  	ps := server.NewPeerServer(psConfig, registry, store, &mb, followersStats, serverStats)
    47  	psListener := server.NewListener("http", testRaftURL, nil)
    48  
    49  	// Create Raft transporter and server
    50  	dialTimeout := (3 * testHeartbeatInterval) + testElectionTimeout
    51  	responseHeaderTimeout := (3 * testHeartbeatInterval) + testElectionTimeout
    52  	raftTransporter := server.NewTransporter(followersStats, serverStats, registry, testHeartbeatInterval, dialTimeout, responseHeaderTimeout)
    53  	raftServer, err := raft.NewServer(testName, path, raftTransporter, store, ps, "")
    54  	if err != nil {
    55  		panic(err)
    56  	}
    57  	raftServer.SetElectionTimeout(testElectionTimeout)
    58  	raftServer.SetHeartbeatInterval(testHeartbeatInterval)
    59  	ps.SetRaftServer(raftServer)
    60  
    61  	s := server.New(testName, "http://"+testClientURL, ps, registry, store, nil)
    62  	sListener := server.NewListener("http", testClientURL, nil)
    63  
    64  	ps.SetServer(s)
    65  
    66  	w := &sync.WaitGroup{}
    67  
    68  	// Start up peer server.
    69  	c := make(chan bool)
    70  	go func() {
    71  		c <- true
    72  		ps.Start(false, "", []string{})
    73  		h := waitHandler{w, ps.HTTPHandler()}
    74  		http.Serve(psListener, &h)
    75  	}()
    76  	<-c
    77  
    78  	// Start up etcd server.
    79  	go func() {
    80  		c <- true
    81  		h := waitHandler{w, s.HTTPHandler()}
    82  		http.Serve(sListener, &h)
    83  	}()
    84  	<-c
    85  
    86  	// Wait to make sure servers have started.
    87  	time.Sleep(50 * time.Millisecond)
    88  
    89  	// Execute the function passed in.
    90  	f(s)
    91  
    92  	// Clean up servers.
    93  	ps.Stop()
    94  	psListener.Close()
    95  	sListener.Close()
    96  	w.Wait()
    97  }
    98  
    99  type waitHandler struct {
   100  	wg      *sync.WaitGroup
   101  	handler http.Handler
   102  }
   103  
   104  func (h *waitHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
   105  	h.wg.Add(1)
   106  	defer h.wg.Done()
   107  	h.handler.ServeHTTP(w, r)
   108  
   109  	//important to flush before decrementing the wait group.
   110  	//we won't get a chance to once main() ends.
   111  	w.(http.Flusher).Flush()
   112  }