chore: wait for the initial store sync and cluster components before serving cluster gRPC requests

internal/machine/cluster.go

@@ -40,7 +40,9 @@ type clusterController struct {
 	corroService corroservice.Service
 	dockerCtrl   *docker.Controller
 	// dockerReady is signalled when Docker is configured and ready for containers.
-	dockerReady     chan<- struct{}
+	dockerReady chan<- struct{}
+	// clusterReady is signalled when the cluster controller has finished initializing all components.
+	clusterReady    chan<- struct{}
 	caddyconfigCtrl *caddyconfig.Controller
 
 	// dnsServer is the embedded internal DNS server for the cluster listening on the machine IP.
@@ -60,6 +62,7 @@ func newClusterController(
 	corroService corroservice.Service,
 	dockerService *docker.Service,
 	dockerReady chan<- struct{},
+	clusterReady chan<- struct{},
 	caddyfileCtrl *caddyconfig.Controller,
 	dnsServer *dns.Server,
 	dnsResolver *dns.ClusterResolver,
@@ -81,6 +84,7 @@ func newClusterController(
 		corroService:    corroService,
 		dockerCtrl:      docker.NewController(state.ID, dockerService, store),
 		dockerReady:     dockerReady,
+		clusterReady:    clusterReady,
 		caddyconfigCtrl: caddyfileCtrl,
 		dnsServer:       dnsServer,
 		dnsResolver:     dnsResolver,
@@ -138,20 +142,6 @@ func (cc *clusterController) Run(ctx context.Context) error {
 		return nil
 	})
 
-	// Wait for the store database to sync to the minimum version before starting store-dependent components.
-	// This prevents issues with using partially replicated data when the machine just joined the cluster,
-	// e.g., an empty machine list causing WireGuard peer misconfiguration.
-	cc.waitStoreSync(ctx)
-
-	// Check if waitStoreSync exited because the context was cancelled. Return early in that case.
-	if ctx.Err() != nil {
-		err := errGroup.Wait()
-		if corroErr := cc.stopCorrosion(); corroErr != nil {
-			err = errors.Join(err, corroErr)
-		}
-		return err
-	}
-
 	// Start the network API server. Assume the management IP can't be changed when the network is running.
 	apiAddr := net.JoinHostPort(cc.state.Network.ManagementIP.String(), strconv.Itoa(constants.MachineAPIPort))
 	listener, err := net.Listen("tcp", apiAddr)
@@ -166,6 +156,22 @@ func (cc *clusterController) Run(ctx context.Context) error {
 		return nil
 	})
 
+	// Wait for the store database to sync to the minimum version before starting store-dependent components.
+	// This prevents issues with using partially replicated data when the machine just joined the cluster,
+	// e.g., an empty machine list causing WireGuard peer misconfiguration.
+	cc.waitStoreSync(ctx)
+
+	// Check if waitStoreSync exited because the context was cancelled. Return early in that case.
+	if ctx.Err() != nil {
+		cc.stopAPIServer()
+
+		err := errGroup.Wait()
+		if corroErr := cc.stopCorrosion(); corroErr != nil {
+			err = errors.Join(err, corroErr)
+		}
+		return err
+	}
+
 	errGroup.Go(func() error {
 		slog.Info("Starting embedded DNS resolver.")
 		if err := cc.dnsResolver.Run(ctx); err != nil {
@@ -216,13 +222,13 @@ func (cc *clusterController) Run(ctx context.Context) error {
 		})
 	}
 
+	// Signal that the cluster controller has finished starting all components.
+	close(cc.clusterReady)
+	slog.Info("Cluster controller finished starting all components.")
 	// Wait for the context to be done and stop all servers and controllers.
 	<-ctx.Done()
 
-	slog.Info("Stopping network API server.")
-	// TODO: implement timeout for graceful shutdown.
-	cc.server.GracefulStop()
-	slog.Info("Network API server stopped.")
+	cc.stopAPIServer()
 
 	// Stop the unregistry server with a timeout if it was started.
 	if cc.unregistry != nil {
@@ -248,6 +254,29 @@ func (cc *clusterController) Run(ctx context.Context) error {
 	return err
 }
 
+// stopAPIServer gracefully stops the network API server with a timeout.
+func (cc *clusterController) stopAPIServer() {
+	timeout := 10 * time.Second
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	stopped := make(chan struct{})
+	go func() {
+		slog.Info("Stopping network API server.")
+		cc.server.GracefulStop()
+		close(stopped)
+	}()
+
+	select {
+	case <-ctx.Done():
+		slog.Warn("Network API server graceful stop timed out, forcing stop.", "timeout", timeout)
+		cc.server.Stop()
+	case <-stopped:
+	}
+
+	slog.Info("Network API server stopped.")
+}
+
 // stopCorrosion stops the Corrosion service with a timeout.
 func (cc *clusterController) stopCorrosion() error {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)

internal/machine/cluster/cluster.go

@@ -26,12 +26,19 @@ type Cluster struct {
 	corroAdmin *corrosion.AdminClient
 	// machineID is the ID of the current machine that is running the cluster service.
 	machineID string
+	// initialised is closed when the machine is configured as a member of a cluster.
+	initialised <-chan struct{}
+	// ready is closed when the cluster controller has finished starting all components
+	// and the machine is ready to serve cluster requests.
+	ready <-chan struct{}
 }
 
-func NewCluster(store *store.Store, corroAdmin *corrosion.AdminClient) *Cluster {
+func NewCluster(store *store.Store, corroAdmin *corrosion.AdminClient, initialised, ready <-chan struct{}) *Cluster {
 	return &Cluster{
-		store:      store,
-		corroAdmin: corroAdmin,
+		store:       store,
+		corroAdmin:  corroAdmin,
+		initialised: initialised,
+		ready:       ready,
 	}
 }
 
@@ -41,67 +48,45 @@ func (c *Cluster) UpdateMachineID(mid string) {
 }
 
 func (c *Cluster) Init(ctx context.Context, network netip.Prefix) error {
-	initialised, err := c.Initialised(ctx)
-	if err != nil {
-		return err
-	}
-	if initialised {
-		return fmt.Errorf("cluster is already initialised")
+	select {
+	case <-c.initialised:
+		return fmt.Errorf("cluster is already initialised on this machine")
+	default:
 	}
 
-	if err = c.store.Put(ctx, "network", network.String()); err != nil {
+	if err := c.store.Put(ctx, "network", network.String()); err != nil {
 		return fmt.Errorf("put network to store: %w", err)
 	}
-	if err = c.store.Put(ctx, "created_at", time.Now().UTC().Format(time.RFC3339)); err != nil {
+	if err := c.store.Put(ctx, "created_at", time.Now().UTC().Format(time.RFC3339)); err != nil {
 		return fmt.Errorf("put created_at to store: %w", err)
 	}
 	return nil
 }
 
-func (c *Cluster) Initialised(ctx context.Context) (bool, error) {
-	var createdAt string
-	if err := c.store.Get(ctx, "created_at", &createdAt); err != nil {
-		if errors.Is(err, store.ErrKeyNotFound) {
-			return false, nil
-		}
-		return false, status.Errorf(codes.Internal, "get created_at from store: %v", err)
+// checkReady checks if the machine is ready to serve cluster requests (store synced, cluster components started).
+func (c *Cluster) checkReady() error {
+	select {
+	case <-c.ready:
+		return nil
+	default:
+		return status.Error(codes.Unavailable, "machine is not ready to serve cluster requests")
 	}
-	return true, nil
-}
-
-func (c *Cluster) checkInitialised(ctx context.Context) error {
-	initialised, err := c.Initialised(ctx)
-	if err != nil {
-		return err
-	}
-	if !initialised {
-		return status.Error(codes.FailedPrecondition, "cluster is not initialised")
-	}
-	return nil
-}
-
-func (c *Cluster) Network(ctx context.Context) (netip.Prefix, error) {
-	if err := c.checkInitialised(ctx); err != nil {
-		return netip.Prefix{}, err
-	}
-
-	var net string
-	if err := c.store.Get(ctx, "network", &net); err != nil {
-		return netip.Prefix{}, status.Errorf(codes.Internal, "get network from store: %v", err)
-	}
-	prefix, err := netip.ParsePrefix(net)
-	if err != nil {
-		return netip.Prefix{}, status.Errorf(codes.Internal, "parse network prefix: %v", err)
-	}
-	return prefix, nil
 }
 
 // AddMachine adds a machine to the cluster.
 func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*pb.AddMachineResponse, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 
+	return c.AddMachineWithoutReadyCheck(ctx, req)
+}
+
+// AddMachineWithoutReadyCheck adds a machine to the cluster without checking if the cluster is ready.
+// This is used internally during cluster initialisation to add the first machine.
+func (c *Cluster) AddMachineWithoutReadyCheck(
+	ctx context.Context, req *pb.AddMachineRequest,
+) (*pb.AddMachineResponse, error) {
 	if req.Network == nil {
 		return nil, status.Error(codes.InvalidArgument, "network not set")
 	}
@@ -162,7 +147,7 @@ func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*p
 		manageIP = pb.NewIP(network.ManagementIP(req.Network.PublicKey))
 	}
 	// Allocate a subnet for the machine from the cluster network.
-	clusterNetwork, err := c.Network(ctx)
+	clusterNetwork, err := c.network(ctx)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "get cluster network: %v", err)
 	}
@@ -198,9 +183,21 @@ func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*p
 	return resp, nil
 }
 
+func (c *Cluster) network(ctx context.Context) (netip.Prefix, error) {
+	var net string
+	if err := c.store.Get(ctx, "network", &net); err != nil {
+		return netip.Prefix{}, status.Errorf(codes.Internal, "get network from store: %v", err)
+	}
+	prefix, err := netip.ParsePrefix(net)
+	if err != nil {
+		return netip.Prefix{}, status.Errorf(codes.Internal, "parse network prefix: %v", err)
+	}
+	return prefix, nil
+}
+
 // UpdateMachine updates machine configuration in the cluster.
 func (c *Cluster) UpdateMachine(ctx context.Context, req *pb.UpdateMachineRequest) (*pb.UpdateMachineResponse, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 
@@ -283,7 +280,7 @@ func (c *Cluster) UpdateMachine(ctx context.Context, req *pb.UpdateMachineReques
 
 // ListMachines lists all machines in the cluster including their membership states.
 func (c *Cluster) ListMachines(ctx context.Context, _ *emptypb.Empty) (*pb.ListMachinesResponse, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 
@@ -329,7 +326,7 @@ func (c *Cluster) ListMachines(ctx context.Context, _ *emptypb.Empty) (*pb.ListM
 
 // RemoveMachine removes a machine from the cluster.
 func (c *Cluster) RemoveMachine(ctx context.Context, req *pb.RemoveMachineRequest) (*emptypb.Empty, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 

internal/machine/cluster/dns.go

@@ -25,7 +25,7 @@ type uncloudDNSDomain struct {
 }
 
 func (c *Cluster) ReserveDomain(ctx context.Context, req *pb.ReserveDomainRequest) (*pb.Domain, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 
@@ -64,7 +64,7 @@ func (c *Cluster) ReserveDomain(ctx context.Context, req *pb.ReserveDomainReques
 }
 
 func (c *Cluster) GetDomain(ctx context.Context, _ *emptypb.Empty) (*pb.Domain, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 
@@ -95,7 +95,7 @@ func (c *Cluster) storedDomain(ctx context.Context) (uncloudDNSDomain, error) {
 }
 
 func (c *Cluster) ReleaseDomain(ctx context.Context, _ *emptypb.Empty) (*pb.Domain, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 
@@ -115,7 +115,7 @@ func (c *Cluster) ReleaseDomain(ctx context.Context, _ *emptypb.Empty) (*pb.Doma
 func (c *Cluster) CreateDomainRecords(
 	ctx context.Context, req *pb.CreateDomainRecordsRequest,
 ) (*pb.CreateDomainRecordsResponse, error) {
-	if err := c.checkInitialised(ctx); err != nil {
+	if err := c.checkReady(); err != nil {
 		return nil, err
 	}
 

internal/machine/machine.go

@@ -174,10 +174,13 @@ type Machine struct {
 	state  *State
 	// started is closed when the machine is ready to serve requests on the local API server.
 	started chan struct{}
-	// initialised is signalled when the machine is configured as a member of a cluster.
+	// initialised is closed when the machine is configured as a member of a cluster.
 	initialised chan struct{}
-	// networkReady is signalled when the Docker network is configured and ready for containers.
+	// networkReady is closed when the Docker network is configured and ready for containers.
 	networkReady chan struct{}
+	// clusterReady is closed when the cluster controller has finished starting all components
+	// and the machine is ready to serve cluster requests.
+	clusterReady chan struct{}
 	// resetting is true when the machine is being reset.
 	resetting bool
 	// stop cancels the Run method context to stop the machine gracefully.
@@ -246,7 +249,10 @@ func NewMachine(config *Config) (*Machine, error) {
 	if err != nil {
 		return nil, fmt.Errorf("create corrosion admin client: %w", err)
 	}
-	c := cluster.NewCluster(corroStore, corroAdmin)
+
+	initialised := make(chan struct{})
+	clusterReady := make(chan struct{})
+	c := cluster.NewCluster(corroStore, corroAdmin, initialised, clusterReady)
 
 	// Init dependencies for a gRPC Docker server that proxies requests to the local Docker daemon.
 	dbFilePath := filepath.Join(config.DataDir, DBFileName)
@@ -269,8 +275,9 @@ func NewMachine(config *Config) (*Machine, error) {
 		config:           *config,
 		state:            state,
 		started:          make(chan struct{}),
-		initialised:      make(chan struct{}, 1),
+		initialised:      initialised,
 		networkReady:     make(chan struct{}),
+		clusterReady:     clusterReady,
 		store:            corroStore,
 		cluster:          c,
 		dockerService:    dockerService,
@@ -294,7 +301,7 @@ func NewMachine(config *Config) (*Machine, error) {
 	m.localMachineServer = newGRPCServer(m, c, m.dockerServer, caddyServer)
 
 	if m.Initialised() {
-		m.initialised <- struct{}{}
+		close(m.initialised)
 	}
 
 	return m, nil
@@ -475,6 +482,7 @@ func (m *Machine) Run(ctx context.Context) error {
 				m.config.CorrosionService,
 				m.dockerService,
 				m.networkReady,
+				m.clusterReady,
 				caddyconfigCtrl,
 				dnsServer,
 				dnsResolver,
@@ -722,7 +730,7 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 		addReq.PublicIp = pb.NewIP(publicIP)
 	}
 
-	addResp, err := m.cluster.AddMachine(ctx, addReq)
+	addResp, err := m.cluster.AddMachineWithoutReadyCheck(ctx, addReq)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "add machine to cluster: %v", err)
 	}
@@ -749,7 +757,7 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 	}
 	slog.Info("Cluster initialised with machine.", "id", m.state.ID, "machine", m.state.Name)
 	// Signal that the machine is initialised as a member of a cluster.
-	m.initialised <- struct{}{}
+	close(m.initialised)
 
 	resp := &pb.InitClusterResponse{
 		Machine: addResp.Machine,
@@ -831,7 +839,7 @@ func (m *Machine) JoinCluster(_ context.Context, req *pb.JoinClusterRequest) (*e
 		"peers", len(m.state.Network.Peers),
 	)
 	// Signal that the machine is initialised as a member of a cluster.
-	m.initialised <- struct{}{}
+	close(m.initialised)
 
 	return &emptypb.Empty{}, nil
 }