func newCluster(tb testing.TB, discoveryClient discovery.Client, servers map[string]*grpc.Server) Cluster {
cluster := cluster{
rolers: make(map[string]role.Roler),
servers: make(map[string]server.CombinedAPIServer),
addresser: route.NewDiscoveryAddresser(
discoveryClient,
testNamespace(),
),
sharder: route.NewSharder(
testShardsPerServer * testNumServers,
),
tb: tb,
}
for address, s := range servers {
combinedAPIServer := server.NewCombinedAPIServer(
cluster.sharder,
route.NewRouter(
cluster.addresser,
grpcutil.NewDialer(
grpc.WithInsecure(),
),
address,
),
getDriver(tb, address),
)
pfs.RegisterApiServer(s, combinedAPIServer)
pfs.RegisterInternalApiServer(s, combinedAPIServer)
roler := role.NewRoler(cluster.addresser, cluster.sharder, combinedAPIServer, address, testNumReplicas)
go func() { require.Equal(tb, roler.Run(), discovery.ErrCancelled) }()
cluster.addresses = append(cluster.addresses, address)
cluster.rolers[address] = roler
cluster.servers[address] = combinedAPIServer
}
return &cluster
}
func runMasterReplicaTest(t *testing.T, client discovery.Client) {
addresser := route.NewDiscoveryAddresser(client, "TestMasterReplicaRoler")
serverGroup1 := NewServerGroup(t, addresser, testNumServers/2, 0, testNumReplicas)
go serverGroup1.run(t)
start := time.Now()
for !serverGroup1.satisfied((testNumShards * (testNumReplicas + 1)) / (testNumServers / 2)) {
time.Sleep(500 * time.Millisecond)
if time.Since(start) > time.Second*time.Duration(30) {
t.Fatal("test timed out")
}
}
serverGroup2 := NewServerGroup(t, addresser, testNumServers/2, testNumServers/2, testNumReplicas)
go serverGroup2.run(t)
start = time.Now()
for !serverGroup1.satisfied((testNumShards*(testNumReplicas+1))/testNumServers) ||
!serverGroup2.satisfied((testNumShards*(testNumReplicas+1))/testNumServers) {
time.Sleep(time.Second)
if time.Since(start) > time.Second*time.Duration(60) {
t.Fatal("test timed out")
}
}
serverGroup1.cancel()
for !serverGroup2.satisfied((testNumShards * (testNumReplicas + 1)) / (testNumServers / 2)) {
time.Sleep(500 * time.Millisecond)
if time.Since(start) > time.Second*time.Duration(60) {
t.Fatal("test timed out")
}
}
}
func registerFunc(driver drive.Driver, discoveryClient discovery.Client, servers map[string]*grpc.Server) {
addresser := route.NewDiscoveryAddresser(
discoveryClient,
testNamespace(),
)
i := 0
for address := range servers {
for j := 0; j < testShardsPerServer; j++ {
// TODO(pedge): error
_ = addresser.SetMasterAddress((i*testShardsPerServer)+j, address, 0)
}
i++
}
for address, s := range servers {
combinedAPIServer := server.NewCombinedAPIServer(
route.NewSharder(
testShardsPerServer*testNumServers,
),
route.NewRouter(
addresser,
grpcutil.NewDialer(),
address,
),
driver,
)
pfs.RegisterApiServer(s, combinedAPIServer)
pfs.RegisterInternalApiServer(s, combinedAPIServer)
}
}
func runTest(
t *testing.T,
driver drive.Driver,
f func(t *testing.T, apiClient pfs.ApiClient),
) {
grpctest.Run(
t,
testNumServers,
func(servers map[string]*grpc.Server) {
discoveryClient := discovery.NewMockClient()
i := 0
addresses := make([]string, testNumServers)
for address := range servers {
shards := make([]string, testShardsPerServer)
for j := 0; j < testShardsPerServer; j++ {
shards[j] = fmt.Sprintf("%d", (i*testShardsPerServer)+j)
}
_ = discoveryClient.Set(address+"-master", strings.Join(shards, ","))
addresses[i] = address
i++
}
_ = discoveryClient.Set("all-addresses", strings.Join(addresses, ","))
for address, server := range servers {
combinedAPIServer := NewCombinedAPIServer(
route.NewSharder(
testShardsPerServer*testNumServers,
),
route.NewRouter(
route.NewDiscoveryAddresser(
discoveryClient,
),
route.NewDialer(),
address,
),
driver,
)
pfs.RegisterApiServer(server, combinedAPIServer)
pfs.RegisterInternalApiServer(server, combinedAPIServer)
}
},
func(t *testing.T, clientConns map[string]*grpc.ClientConn) {
var clientConn *grpc.ClientConn
for _, c := range clientConns {
clientConn = c
break
}
for _, c := range clientConns {
if c != clientConn {
_ = c.Close()
}
}
f(
t,
pfs.NewApiClient(
clientConn,
),
)
},
)
}
func runTest(t *testing.T, client discovery.Client) {
addresser := route.NewDiscoveryAddresser(client, "TestRoler")
serverGroup1 := NewServerGroup(addresser, testNumServers/2, 0)
go serverGroup1.run(t)
start := time.Now()
for !serverGroup1.satisfied(testNumShards / (testNumServers / 2)) {
time.Sleep(3 * time.Second)
if time.Since(start) > time.Second*time.Duration(10) {
t.Fatal("test timed out")
}
}
serverGroup2 := NewServerGroup(addresser, testNumServers/2, testNumServers/2)
go serverGroup2.run(t)
start = time.Now()
for !serverGroup1.satisfied(testNumShards/testNumServers) || !serverGroup2.satisfied(testNumShards/testNumServers) {
time.Sleep(3 * time.Second)
if time.Since(start) > time.Second*time.Duration(10) {
t.Fatal("test timed out")
}
}
serverGroup1.cancel()
for !serverGroup2.satisfied(testNumShards) {
time.Sleep(3 * time.Second)
if time.Since(start) > time.Second*time.Duration(30) {
t.Fatal("test timed out")
}
}
}
func do(appEnvObj interface{}) error {
appEnv := appEnvObj.(*appEnv)
discoveryClient, err := getEtcdClient()
if err != nil {
return err
}
address := fmt.Sprintf("0.0.0.0:%d", appEnv.APIPort)
addresser := route.NewDiscoveryAddresser(
discoveryClient,
"namespace",
)
for i := 0; i < appEnv.NumShards; i++ {
if err := addresser.SetMasterAddress(i, address, 0); err != nil {
return err
}
}
var driver drive.Driver
switch appEnv.DriverType {
case "btrfs":
driver, err = btrfs.NewDriver(appEnv.DriverRoot, "")
if err != nil {
return err
}
default:
return fmt.Errorf("unknown value for PFS_DRIVER_TYPE: %s", appEnv.DriverType)
}
combinedAPIServer := server.NewCombinedAPIServer(
route.NewSharder(
appEnv.NumShards,
),
route.NewRouter(
addresser,
grpcutil.NewDialer(),
address,
),
driver,
)
return grpcutil.GrpcDo(
appEnv.APIPort,
appEnv.TracePort,
pachyderm.Version,
func(s *grpc.Server) {
pfs.RegisterApiServer(s, combinedAPIServer)
pfs.RegisterInternalApiServer(s, combinedAPIServer)
},
)
}
func registerFunc(tb testing.TB, discoveryClient discovery.Client, servers map[string]*grpc.Server) error {
addresser := route.NewDiscoveryAddresser(
discoveryClient,
testNamespace(),
)
i := 0
for address := range servers {
for j := 0; j < testShardsPerServer; j++ {
if err := addresser.SetMasterAddress((i*testShardsPerServer)+j, address, 0); err != nil {
return err
}
if err := addresser.SetReplicaAddress((((i+1)%len(servers))*testShardsPerServer)+j, address, 0); err != nil {
return err
}
if err := addresser.SetReplicaAddress((((i+2)%len(servers))*testShardsPerServer)+j, address, 0); err != nil {
return err
}
}
i++
}
for address, s := range servers {
combinedAPIServer := server.NewCombinedAPIServer(
route.NewSharder(
testShardsPerServer*testNumServers,
),
route.NewRouter(
addresser,
grpcutil.NewDialer(),
address,
),
getDriver(tb, address),
)
pfs.RegisterApiServer(s, combinedAPIServer)
pfs.RegisterInternalApiServer(s, combinedAPIServer)
}
return nil
}
func do(appEnvObj interface{}) error {
appEnv := appEnvObj.(*appEnv)
logrus.Register()
discoveryClient, err := getEtcdClient()
if err != nil {
return err
}
address := appEnv.Address
if address == "" {
address, err = netutil.ExternalIP()
if err != nil {
return err
}
}
sharder := route.NewSharder(appEnv.NumShards, 0)
address = fmt.Sprintf("%s:%d", address, appEnv.Port)
addresser := route.NewDiscoveryAddresser(
discoveryClient,
sharder,
"namespace",
)
var driver drive.Driver
switch appEnv.DriverType {
case "btrfs":
driver, err = btrfs.NewDriver(appEnv.DriverRoot, "")
if err != nil {
return err
}
default:
return fmt.Errorf("unknown value for PFS_DRIVER_TYPE: %s", appEnv.DriverType)
}
apiServer := server.NewAPIServer(
route.NewSharder(
appEnv.NumShards,
0,
),
route.NewRouter(
addresser,
grpcutil.NewDialer(
grpc.WithInsecure(),
),
address,
),
)
internalAPIServer := server.NewInternalAPIServer(
route.NewSharder(
appEnv.NumShards,
0,
),
route.NewRouter(
addresser,
grpcutil.NewDialer(
grpc.WithInsecure(),
),
address,
),
driver,
)
go func() {
if err := addresser.Register(nil, "id", address, internalAPIServer); err != nil {
log.Print(err)
}
}()
go func() {
if err := addresser.AssignRoles(nil); err != nil {
log.Print(err)
}
}()
// TODO(pedge): no!
trace.AuthRequest = func(_ *http.Request) (bool, bool) {
return true, true
}
return protoserver.Serve(
uint16(appEnv.Port),
func(s *grpc.Server) {
pfs.RegisterApiServer(s, apiServer)
pfs.RegisterInternalApiServer(s, internalAPIServer)
},
protoserver.ServeOptions{
HTTPPort: uint16(appEnv.HTTPPort),
DebugPort: uint16(appEnv.DebugPort),
Version: pachyderm.Version,
HTTPRegisterFunc: func(ctx context.Context, mux *runtime.ServeMux, clientConn *grpc.ClientConn) error {
return pfs.RegisterApiHandler(ctx, mux, clientConn)
},
},
)
}
func do(appEnvObj interface{}) error {
appEnv := appEnvObj.(*appEnv)
discoveryClient, err := getEtcdClient()
if err != nil {
return err
}
address := appEnv.Address
if address == "" {
address, err = netutil.ExternalIP()
if err != nil {
return err
}
}
address = fmt.Sprintf("%s:%d", address, appEnv.Port)
addresser := route.NewDiscoveryAddresser(
discoveryClient,
"namespace",
)
for i := 0; i < appEnv.NumShards; i++ {
if _, err := addresser.SetMasterAddress(i, route.Address{address, false}); err != nil {
return err
}
}
var driver drive.Driver
switch appEnv.DriverType {
case "btrfs":
driver, err = btrfs.NewDriver(appEnv.DriverRoot, "")
if err != nil {
return err
}
default:
return fmt.Errorf("unknown value for PFS_DRIVER_TYPE: %s", appEnv.DriverType)
}
combinedAPIServer := server.NewCombinedAPIServer(
route.NewSharder(
appEnv.NumShards,
),
route.NewRouter(
addresser,
grpcutil.NewDialer(
grpc.WithInsecure(),
),
address,
),
driver,
)
return protoserver.Serve(
uint16(appEnv.Port),
func(s *grpc.Server) {
pfs.RegisterApiServer(s, combinedAPIServer)
pfs.RegisterInternalApiServer(s, combinedAPIServer)
},
protoserver.ServeOptions{
HTTPPort: uint16(appEnv.HTTPPort),
TracePort: uint16(appEnv.TracePort),
Version: pachyderm.Version,
HTTPRegisterFunc: func(ctx context.Context, mux *runtime.ServeMux, clientConn *grpc.ClientConn) error {
return pfs.RegisterApiHandler(ctx, mux, clientConn)
},
},
)
}