func RunBench(
	b *testing.B,
	f func(*testing.B, pfs.ApiClient),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(b, err)
	var cluster Cluster
	prototest.RunB(
		b,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			cluster = registerFunc(b, discoveryClient, servers)
		},
		func(b *testing.B, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			f(
				b,
				pfs.NewApiClient(
					clientConn,
				),
			)
		},
	)
	cluster.Shutdown()
}

func (a *combinedAPIServer) GetFile(getFileRequest *pfs.GetFileRequest, apiGetFileServer pfs.Api_GetFileServer) (retErr error) {
	shard, clientConn, err := a.getShardAndClientConnIfNecessary(getFileRequest.Path, true)
	if err != nil {
		return err
	}
	if clientConn != nil {
		apiGetFileClient, err := pfs.NewApiClient(clientConn).GetFile(context.Background(), getFileRequest)
		if err != nil {
			return err
		}
		return protoutil.RelayFromStreamingBytesClient(apiGetFileClient, apiGetFileServer)
	}
	file, err := a.driver.GetFile(getFileRequest.Path, shard)
	if err != nil {
		return err
	}
	defer func() {
		if err := file.Close(); err != nil && retErr == nil {
			retErr = err
		}
	}()
	return protoutil.WriteToStreamingBytesServer(
		io.NewSectionReader(file, getFileRequest.OffsetBytes, getFileRequest.SizeBytes),
		apiGetFileServer,
	)
}

func (a *combinedAPIServer) MakeDirectory(ctx context.Context, makeDirectoryRequest *pfs.MakeDirectoryRequest) (*google_protobuf.Empty, error) {
	shards, err := a.getAllShards(true)
	if err != nil {
		return nil, err
	}
	if err := a.driver.MakeDirectory(makeDirectoryRequest.Path, shards); err != nil {
		return nil, err
	}
	if !makeDirectoryRequest.Redirect {
		clientConns, err := a.router.GetAllClientConns()
		if err != nil {
			return nil, err
		}
		for _, clientConn := range clientConns {
			if _, err := pfs.NewApiClient(clientConn).MakeDirectory(
				ctx,
				&pfs.MakeDirectoryRequest{
					Path:     makeDirectoryRequest.Path,
					Redirect: true,
				},
			); err != nil {
				return nil, err
			}
		}
	}
	return emptyInstance, nil
}

func RunBench(
	b *testing.B,
	f func(b *testing.B, apiClient pfs.ApiClient),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(b, err)
	grpctest.RunB(
		b,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			registerFunc(getDriver(b), discoveryClient, servers)
		},
		func(b *testing.B, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			f(
				b,
				pfs.NewApiClient(
					clientConn,
				),
			)
		},
	)
}

func RunTest(
	t *testing.T,
	f func(*testing.T, pfs.ApiClient, pfs.InternalApiClient, Cluster),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(t, err)
	var cluster Cluster
	prototest.RunT(
		t,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			cluster = registerFunc(t, discoveryClient, servers)
		},
		func(t *testing.T, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			f(
				t,
				pfs.NewApiClient(
					clientConn,
				),
				pfs.NewInternalApiClient(
					clientConn,
				),
				cluster,
			)
		},
	)
	cluster.Shutdown()
}

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 (a *combinedAPIServer) Commit(ctx context.Context, commitRequest *pfs.CommitRequest) (*google_protobuf.Empty, error) {
	shards, err := a.getAllShards(false)
	if err != nil {
		return nil, err
	}
	if err := a.driver.Commit(commitRequest.Commit, shards); err != nil {
		return nil, err
	}
	if !commitRequest.Redirect {
		clientConns, err := a.router.GetAllClientConns()
		if err != nil {
			return nil, err
		}
		for _, clientConn := range clientConns {
			if _, err := pfs.NewApiClient(clientConn).Commit(
				ctx,
				&pfs.CommitRequest{
					Commit:   commitRequest.Commit,
					Redirect: true,
				},
			); err != nil {
				return nil, err
			}
		}
	}
	return emptyInstance, nil
}

func RunTest(
	t *testing.T,
	f func(t *testing.T, apiClient pfs.ApiClient, internalAPIClient pfs.InternalApiClient),
) {
	discoveryClient, err := getEtcdClient()
	require.NoError(t, err)
	grpctest.Run(
		t,
		testNumServers,
		func(servers map[string]*grpc.Server) {
			registerFunc(getDriver(t), discoveryClient, servers)
		},
		func(t *testing.T, clientConns map[string]*grpc.ClientConn) {
			var clientConn *grpc.ClientConn
			for _, c := range clientConns {
				clientConn = c
				break
			}
			f(
				t,
				pfs.NewApiClient(
					clientConn,
				),
				pfs.NewInternalApiClient(
					clientConn,
				),
			)
		},
	)
}

func do(appEnvObj interface{}) error {
	appEnv := appEnvObj.(*appEnv)
	logrus.Register()
	address := appEnv.PachydermPfsd1Port
	if address == "" {
		address = appEnv.PfsAddress
	} else {
		address = strings.Replace(address, "tcp://", "", -1)
	}
	return dockervolume.Serve(
		dockervolume.NewVolumeDriverHandler(
			newVolumeDriver(
				func() (fuse.Mounter, error) {
					clientConn, err := grpc.Dial(address, grpc.WithInsecure())
					if err != nil {
						return nil, err
					}
					return fuse.NewMounter(
						pfs.NewApiClient(
							clientConn,
						),
					), nil
				},
				appEnv.BaseMountpoint,
			),
			dockervolume.VolumeDriverHandlerOptions{},
		),
		dockervolume.ProtocolUnix,
		volumeDriverName,
		volumeDriverGroup,
	)
}

func (a *combinedAPIServer) InitRepository(ctx context.Context, initRepositoryRequest *pfs.InitRepositoryRequest) (*google_protobuf.Empty, error) {
	shards, err := a.getAllShards(true)
	if err != nil {
		return nil, err
	}
	if err := a.driver.InitRepository(initRepositoryRequest.Repository, initRepositoryRequest.Replica, shards); err != nil {
		return nil, err
	}
	if !initRepositoryRequest.Redirect {
		clientConns, err := a.router.GetAllClientConns()
		if err != nil {
			return nil, err
		}
		for _, clientConn := range clientConns {
			if _, err := pfs.NewApiClient(clientConn).InitRepository(
				ctx,
				&pfs.InitRepositoryRequest{
					Repository: initRepositoryRequest.Repository,
					Replica:    initRepositoryRequest.Replica,
					Redirect:   true,
				},
			); err != nil {
				return nil, err
			}
		}
	}
	return emptyInstance, nil
}

func (a *combinedAPIServer) Branch(ctx context.Context, branchRequest *pfs.BranchRequest) (*pfs.BranchResponse, error) {
	if branchRequest.Redirect && branchRequest.NewCommit == nil {
		return nil, fmt.Errorf("must set a new commit for redirect %+v", branchRequest)
	}
	shards, err := a.getAllShards(false)
	if err != nil {
		return nil, err
	}
	newCommit, err := a.driver.Branch(branchRequest.Commit, branchRequest.NewCommit, shards)
	if err != nil {
		return nil, err
	}
	if !branchRequest.Redirect {
		clientConns, err := a.router.GetAllClientConns()
		if err != nil {
			return nil, err
		}
		for _, clientConn := range clientConns {
			if _, err := pfs.NewApiClient(clientConn).Branch(
				ctx,
				&pfs.BranchRequest{
					Commit:    branchRequest.Commit,
					Redirect:  true,
					NewCommit: newCommit,
				},
			); err != nil {
				return nil, err
			}
		}
	}
	return &pfs.BranchResponse{
		Commit: newCommit,
	}, nil
}

func (m *mounterProvider) getOnce() (Mounter, error) {
	clientConn, err := grpc.Dial(m.pfsAddress, grpc.WithInsecure())
	if err != nil {
		return nil, err
	}
	return NewMounter(
		pfs.NewApiClient(
			clientConn,
		),
	), nil
}

func (a *combinedAPIServer) InitRepository(ctx context.Context, initRepositoryRequest *pfs.InitRepositoryRequest) (*google_protobuf.Empty, error) {
	masterShards, err := a.router.GetMasterShards()
	if err != nil {
		return nil, err
	}
	if err := a.driver.InitRepository(initRepositoryRequest.Repository, masterShards); err != nil {
		return nil, err
	}
	replicaShards, err := a.router.GetReplicaShards()
	if err != nil {
		return nil, err
	}
	if err := a.driver.InitRepository(initRepositoryRequest.Repository, replicaShards); err != nil {
		return nil, err
	}
	if !initRepositoryRequest.Redirect {
		clientConns, err := a.router.GetAllClientConns()
		if err != nil {
			return nil, err
		}
		for _, clientConn := range clientConns {
			if _, err := pfs.NewApiClient(clientConn).InitRepository(
				ctx,
				&pfs.InitRepositoryRequest{
					Repository: initRepositoryRequest.Repository,
					Redirect:   true,
				},
			); err != nil {
				return nil, err
			}
		}
		// Create the initial commit
		if _, err = a.Branch(ctx, &pfs.BranchRequest{
			Commit: nil,
			NewCommit: &pfs.Commit{
				Repository: initRepositoryRequest.Repository,
				Id:         InitialCommitID,
			},
			Redirect: false,
		}); err != nil {
			return nil, err
		}
		if _, err = a.Commit(ctx, &pfs.CommitRequest{
			Commit: &pfs.Commit{
				Repository: initRepositoryRequest.Repository,
				Id:         InitialCommitID,
			},
			Redirect: false,
		}); err != nil {
			return nil, err
		}
	}
	return emptyInstance, nil
}

func (a *combinedAPIServer) PutFile(ctx context.Context, putFileRequest *pfs.PutFileRequest) (*google_protobuf.Empty, error) {
	shard, clientConn, err := a.getShardAndClientConnIfNecessary(putFileRequest.Path, false)
	if err != nil {
		return nil, err
	}
	if clientConn != nil {
		return pfs.NewApiClient(clientConn).PutFile(ctx, putFileRequest)
	}
	if err := a.driver.PutFile(putFileRequest.Path, shard, bytes.NewReader(putFileRequest.Value)); err != nil {
		return nil, err
	}
	return emptyInstance, nil
}

func getPfsAPIClient(address string) (pfs.ApiClient, error) {
	var err error
	if address == "" {
		address, err = getPfsAddress()
		if err != nil {
			return nil, err
		}
	}
	clientConn, err := grpc.Dial(address)
	if err != nil {
		return nil, err
	}
	return pfs.NewApiClient(clientConn), nil
}

func (a *combinedAPIServer) Master(shard int) error {
	clientConns, err := a.router.GetReplicaClientConns(shard)
	if err != nil {
		return err
	}
	for _, clientConn := range clientConns {
		apiClient := pfs.NewApiClient(clientConn)
		response, err := apiClient.ListRepositories(context.Background(), &pfs.ListRepositoriesRequest{})
		if err != nil {
			return err
		}
		for _, repository := range response.Repository {
			if err := a.driver.InitRepository(repository, map[int]bool{shard: true}); err != nil {
				return err
			}
			response, err := apiClient.ListCommits(context.Background(), &pfs.ListCommitsRequest{Repository: repository})
			if err != nil {
				return err
			}
			localCommitInfo, err := a.driver.ListCommits(repository, shard)
			if err != nil {
				return err
			}
			for i, commitInfo := range response.CommitInfo {
				if i < len(localCommitInfo) {
					if *commitInfo != *localCommitInfo[i] {
						return fmt.Errorf("divergent data")
					}
					continue
				}
				pullDiffClient, err := pfs.NewInternalApiClient(clientConn).PullDiff(
					context.Background(),
					&pfs.PullDiffRequest{
						Commit: commitInfo.Commit,
						Shard:  uint64(shard),
					},
				)
				if err != nil {
					return err
				}
				if err := a.driver.PushDiff(commitInfo.Commit, protostream.NewStreamingBytesReader(pullDiffClient)); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

func (a *apiServer) DeleteCommit(ctx context.Context, request *pfs.DeleteCommitRequest) (*google_protobuf.Empty, error) {
	version, ctx, err := a.versionAndCtx(ctx)
	if err != nil {
		return nil, err
	}
	clientConns, err := a.router.GetAllClientConns(version)
	if err != nil {
		return nil, err
	}
	for _, clientConn := range clientConns {
		if _, err := pfs.NewApiClient(clientConn).DeleteCommit(ctx, request); err != nil {
			return nil, err
		}
	}
	return emptyInstance, nil
}

// TODO(pedge): race on Branch
func (a *combinedAPIServer) GetCommitInfo(ctx context.Context, getCommitInfoRequest *pfs.GetCommitInfoRequest) (*pfs.GetCommitInfoResponse, error) {
	shard, clientConn, err := a.getMasterShardOrMasterClientConnIfNecessary()
	if err != nil {
		return nil, err
	}
	if clientConn != nil {
		return pfs.NewApiClient(clientConn).GetCommitInfo(ctx, getCommitInfoRequest)
	}
	commitInfo, err := a.driver.GetCommitInfo(getCommitInfoRequest.Commit, shard)
	if err != nil {
		return nil, err
	}
	return &pfs.GetCommitInfoResponse{
		CommitInfo: commitInfo,
	}, nil
}

func (a *combinedAPIServer) ListCommits(ctx context.Context, listCommitsRequest *pfs.ListCommitsRequest) (*pfs.ListCommitsResponse, error) {
	shard, clientConn, err := a.getMasterShardOrMasterClientConnIfNecessary()
	if err != nil {
		return nil, err
	}
	if clientConn != nil {
		return pfs.NewApiClient(clientConn).ListCommits(ctx, listCommitsRequest)
	}
	commitInfos, err := a.driver.ListCommits(listCommitsRequest.Repository, shard)
	if err != nil {
		return nil, err
	}
	return &pfs.ListCommitsResponse{
		CommitInfo: commitInfos,
	}, nil
}

func (a *combinedAPIServer) ListFiles(ctx context.Context, listFilesRequest *pfs.ListFilesRequest) (*pfs.ListFilesResponse, error) {
	shards, err := a.getAllShards(false)
	if err != nil {
		return nil, err
	}
	dynamicShard := listFilesRequest.Shard
	if dynamicShard == nil {
		dynamicShard = &pfs.Shard{Number: 0, Modulo: 1}
	}
	filteredShards := make(map[int]bool)
	for shard := range shards {
		if uint64(shard)%dynamicShard.Modulo == dynamicShard.Number {
			filteredShards[shard] = true
		}
	}
	var fileInfos []*pfs.FileInfo
	for shard := range filteredShards {
		subFileInfos, err := a.driver.ListFiles(listFilesRequest.Path, shard)
		if err != nil {
			return nil, err
		}
		fileInfos = append(fileInfos, subFileInfos...)
	}
	if !listFilesRequest.Redirect {
		clientConns, err := a.router.GetAllClientConns()
		if err != nil {
			return nil, err
		}
		for _, clientConn := range clientConns {
			listFilesResponse, err := pfs.NewApiClient(clientConn).ListFiles(
				ctx,
				&pfs.ListFilesRequest{
					Path:     listFilesRequest.Path,
					Shard:    listFilesRequest.Shard,
					Redirect: true,
				},
			)
			if err != nil {
				return nil, err
			}
			fileInfos = append(fileInfos, listFilesResponse.FileInfo...)
		}
	}
	return &pfs.ListFilesResponse{
		FileInfo: fileInfos,
	}, nil
}