seaweedfs/weed/shell/commands_test.go

package shell

import (
	"context"
	"fmt"
	"net"
	"sync/atomic"
	"testing"
	"time"

	"github.com/seaweedfs/seaweedfs/weed/pb/master_pb"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/status"
)

type fakeDiscoveryServer struct {
	master_pb.UnimplementedSeaweedServer
	leader  string
	filers  atomic.Int32
	listErr atomic.Value
}

func (s *fakeDiscoveryServer) KeepConnected(stream master_pb.Seaweed_KeepConnectedServer) error {
	if _, err := stream.Recv(); err != nil {
		return err
	}
	if err := stream.Send(&master_pb.KeepConnectedResponse{
		VolumeLocation: &master_pb.VolumeLocation{Leader: s.leader},
	}); err != nil {
		return err
	}
	<-stream.Context().Done()
	return nil
}

func (s *fakeDiscoveryServer) ListClusterNodes(context.Context, *master_pb.ListClusterNodesRequest) (*master_pb.ListClusterNodesResponse, error) {
	if err, _ := s.listErr.Load().(error); err != nil {
		return nil, err
	}
	resp := &master_pb.ListClusterNodesResponse{}
	for i := int32(0); i < s.filers.Load(); i++ {
		resp.ClusterNodes = append(resp.ClusterNodes, &master_pb.ListClusterNodesResponse_ClusterNode{
			Address: fmt.Sprintf("127.0.0.1:88%02d", i),
		})
	}
	return resp, nil
}

func startFakeDiscoveryMaster(t *testing.T, srv *fakeDiscoveryServer) string {
	t.Helper()
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	grpcServer := grpc.NewServer()
	master_pb.RegisterSeaweedServer(grpcServer, srv)
	go func() { _ = grpcServer.Serve(lis) }()
	t.Cleanup(grpcServer.GracefulStop)

	_, port, _ := net.SplitHostPort(lis.Addr().String())
	return fmt.Sprintf("127.0.0.1:0.%s", port)
}

func discoveryEnv(addr string, timeout time.Duration) (*CommandEnv, context.CancelFunc) {
	masters := addr
	filerGroup := ""
	options := &ShellOptions{
		Masters:        &masters,
		GrpcDialOption: grpc.WithTransportCredentials(insecure.NewCredentials()),
		FilerGroup:     &filerGroup,
		Directory:      "/",
		ConnectTimeout: timeout,
	}
	ce := NewCommandEnv(options)
	ctx, cancel := context.WithCancel(context.Background())
	go ce.MasterClient.KeepConnectedToMaster(ctx)
	ce.MasterClient.WaitUntilConnected(ctx)
	return ce, cancel
}

func TestFilerDiscoveryRetriesUntilFilerAppears(t *testing.T) {
	srv := &fakeDiscoveryServer{}
	addr := startFakeDiscoveryMaster(t, srv)
	srv.leader = addr

	ce, cancel := discoveryEnv(addr, 10*time.Second)
	defer cancel()

	go func() {
		time.Sleep(1200 * time.Millisecond)
		srv.filers.Store(1)
	}()

	got := ce.FilerAddress()
	if got != "127.0.0.1:8800" {
		t.Fatalf("expected the late-registering filer, got %q", got)
	}
}

func TestFilerDiscoveryReturnsEmptyAtDeadline(t *testing.T) {
	srv := &fakeDiscoveryServer{}
	addr := startFakeDiscoveryMaster(t, srv)
	srv.leader = addr

	ce, cancel := discoveryEnv(addr, 1500*time.Millisecond)
	defer cancel()

	start := time.Now()
	if got := ce.FilerAddress(); got != "" {
		t.Fatalf("expected no filer, got %q", got)
	}
	if elapsed := time.Since(start); elapsed < time.Second {
		t.Fatalf("discovery returned after %v, expected it to poll until the deadline", elapsed)
	}
}

func TestFilerDiscoveryRetriesAfterEmptyResult(t *testing.T) {
	srv := &fakeDiscoveryServer{}
	addr := startFakeDiscoveryMaster(t, srv)
	srv.leader = addr

	ce, cancel := discoveryEnv(addr, 1500*time.Millisecond)
	defer cancel()

	if got := ce.FilerAddress(); got != "" {
		t.Fatalf("expected no filer, got %q", got)
	}
	srv.filers.Store(1)
	if got := ce.FilerAddress(); got != "127.0.0.1:8800" {
		t.Fatalf("expected the late-registered filer on retry, got %q", got)
	}
}

func TestFilerDiscoveryStopsOnNonTransientError(t *testing.T) {
	srv := &fakeDiscoveryServer{}
	addr := startFakeDiscoveryMaster(t, srv)
	srv.leader = addr
	srv.listErr.Store(status.Error(codes.PermissionDenied, "denied"))

	ce, cancel := discoveryEnv(addr, 30*time.Second)
	defer cancel()

	start := time.Now()
	if got := ce.FilerAddress(); got != "" {
		t.Fatalf("expected no filer, got %q", got)
	}
	if elapsed := time.Since(start); elapsed > 10*time.Second {
		t.Fatalf("discovery kept polling a non-transient error for %v", elapsed)
	}
}

func TestFilerAddressConfiguredSkipsDiscovery(t *testing.T) {
	ce := &CommandEnv{option: &ShellOptions{FilerAddress: "127.0.0.1:8888"}}
	if got := ce.FilerAddress(); got != "127.0.0.1:8888" {
		t.Fatalf("expected the configured filer, got %q", got)
	}
}