github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/libraries/ngaut/go-zookeeper/zk/server_help.go (about)

     1  package zk
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  	"math/rand"
     7  	"os"
     8  	"path/filepath"
     9  	"time"
    10  )
    11  
    12  type TestServer struct {
    13  	Port int
    14  	Path string
    15  	Srv  *Server
    16  }
    17  
    18  type TestCluster struct {
    19  	Path    string
    20  	Servers []TestServer
    21  }
    22  
    23  func StartTestCluster(size int) (*TestCluster, error) {
    24  	tmpPath, err := ioutil.TempDir("", "gozk")
    25  	if err != nil {
    26  		return nil, err
    27  	}
    28  	success := false
    29  	startPort := int(rand.Int31n(6000) + 10000)
    30  	cluster := &TestCluster{Path: tmpPath}
    31  	defer func() {
    32  		if !success {
    33  			cluster.Stop()
    34  		}
    35  	}()
    36  	for serverN := 0; serverN < size; serverN++ {
    37  		srvPath := filepath.Join(tmpPath, fmt.Sprintf("srv%d", serverN))
    38  		if err := os.Mkdir(srvPath, 0700); err != nil {
    39  			return nil, err
    40  		}
    41  		port := startPort + serverN*3
    42  		cfg := ServerConfig{
    43  			ClientPort: port,
    44  			DataDir:    srvPath,
    45  		}
    46  		for i := 0; i < size; i++ {
    47  			cfg.Servers = append(cfg.Servers, ServerConfigServer{
    48  				ID:                 i + 1,
    49  				Host:               "127.0.0.1",
    50  				PeerPort:           startPort + i*3 + 1,
    51  				LeaderElectionPort: startPort + i*3 + 2,
    52  			})
    53  		}
    54  		cfgPath := filepath.Join(srvPath, "zoo.cfg")
    55  		fi, err := os.Create(cfgPath)
    56  		if err != nil {
    57  			return nil, err
    58  		}
    59  		err = cfg.Marshall(fi)
    60  		fi.Close()
    61  		if err != nil {
    62  			return nil, err
    63  		}
    64  
    65  		fi, err = os.Create(filepath.Join(srvPath, "myid"))
    66  		if err != nil {
    67  			return nil, err
    68  		}
    69  		_, err = fmt.Fprintf(fi, "%d\n", serverN+1)
    70  		fi.Close()
    71  		if err != nil {
    72  			return nil, err
    73  		}
    74  
    75  		srv := &Server{
    76  			ConfigPath: cfgPath,
    77  		}
    78  		if err := srv.Start(); err != nil {
    79  			fmt.Println(err)
    80  			return nil, err
    81  		}
    82  		cluster.Servers = append(cluster.Servers, TestServer{
    83  			Path: srvPath,
    84  			Port: cfg.ClientPort,
    85  			Srv:  srv,
    86  		})
    87  	}
    88  	success = true
    89  	time.Sleep(time.Second) // Give the server time to become active. Should probably actually attempt to connect to verify.
    90  	return cluster, nil
    91  }
    92  
    93  func (ts *TestCluster) Connect(idx int) (*Conn, error) {
    94  	zk, _, err := Connect([]string{fmt.Sprintf("127.0.0.1:%d", ts.Servers[idx].Port)}, time.Second*15)
    95  	return zk, err
    96  }
    97  
    98  func (ts *TestCluster) ConnectAll() (*Conn, error) {
    99  	hosts := make([]string, len(ts.Servers))
   100  	for i, srv := range ts.Servers {
   101  		hosts[i] = fmt.Sprintf("127.0.0.1:%d", srv.Port)
   102  	}
   103  	zk, _, err := Connect(hosts, time.Second*15)
   104  	return zk, err
   105  }
   106  
   107  func (ts *TestCluster) Stop() error {
   108  	for _, srv := range ts.Servers {
   109  		srv.Srv.Stop()
   110  	}
   111  	defer os.RemoveAll(ts.Path)
   112  	return nil
   113  }