refactor cluster server to initialise state when initialising cluster, allow machine to be reinitialised

internal/machine/cluster/server.go

@@ -12,26 +12,38 @@ import (
 	"uncloud/internal/machine/network"
 )
 
-type Server struct {
+type Cluster struct {
 	pb.UnimplementedClusterServer
 
 	state *State
 }
 
-func NewServer(state *State) *Server {
+func NewCluster(state *State) *Cluster {
-	return &Server{
+	return &Cluster{
 		state: state,
 	}
 }
 
-func (c *Server) Network() (netip.Prefix, error) {
+func (c *Cluster) SetState(state *State) {
+	c.state = state
+}
+
+func (c *Cluster) Network() (netip.Prefix, error) {
+	if c.state == nil {
+		return netip.Prefix{}, status.Error(codes.FailedPrecondition, "cluster is not initialized")
+	}
+
 	if c.state.State.Network == nil {
 		return netip.Prefix{}, fmt.Errorf("network not set")
 	}
 	return c.state.State.Network.ToPrefix()
 }
 
-func (c *Server) SetNetwork(network *pb.IPPrefix) error {
+func (c *Cluster) SetNetwork(network *pb.IPPrefix) error {
+	if c.state == nil {
+		return status.Error(codes.FailedPrecondition, "cluster is not initialized")
+	}
+
 	if c.state.State.Network != nil {
 		return fmt.Errorf("network already set and cannot be changed")
 	}
@@ -43,7 +55,11 @@ func (c *Server) SetNetwork(network *pb.IPPrefix) error {
 }
 
 // AddMachine adds a machine to the cluster.
-func (c *Server) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*pb.AddMachineResponse, error) {
+func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*pb.AddMachineResponse, error) {
+	if c.state == nil {
+		return nil, status.Error(codes.FailedPrecondition, "cluster is not initialized")
+	}
+
 	// TODO: replace errors with gRPC status.Error(f), e.g. status.Error(codes.InvalidArgument, "management IP not set")
 	if req.Network.PublicKey == nil {
 		return nil, fmt.Errorf("public key not set")
@@ -126,7 +142,13 @@ func (c *Server) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*pb
 	return resp, nil
 }
 
-func (c *Server) ListMachineEndpoints(ctx context.Context, req *pb.ListMachineEndpointsRequest) (*pb.ListMachineEndpointsResponse, error) {
+func (c *Cluster) ListMachineEndpoints(
+	ctx context.Context, req *pb.ListMachineEndpointsRequest,
+) (*pb.ListMachineEndpointsResponse, error) {
+	if c.state == nil {
+		return nil, status.Error(codes.FailedPrecondition, "cluster is not initialized")
+	}
+
 	endpoints, ok := c.state.State.Endpoints[req.Id]
 	if !ok {
 		return nil, status.Errorf(codes.NotFound, "machine %q not found", req.Id)
@@ -134,16 +156,11 @@ func (c *Server) ListMachineEndpoints(ctx context.Context, req *pb.ListMachineEn
 	return &pb.ListMachineEndpointsResponse{Endpoints: endpoints}, nil
 }
 
-func (c *Server) AddUser(user *pb.User) error {
+func (c *Cluster) AddUser(user *pb.User) error {
 	c.state.State.Users = append(c.state.State.Users, user)
 	return c.state.Save()
 }
 
-func (c *Server) ListUsers() []*pb.User {
+func (c *Cluster) ListUsers() []*pb.User {
 	return c.state.State.Users
 }
-
-type State struct {
-	State *pb.State
-	path  string
-}

internal/machine/cluster/state.go

@@ -10,6 +10,11 @@ import (
 
 const StateFile = "cluster.pb"
 
+type State struct {
+	State *pb.State
+	path  string
+}
+
 func StatePath(dataDir string) string {
 	return filepath.Join(dataDir, StateFile)
 }

internal/machine/machine.go

@@ -36,14 +36,15 @@ type Machine struct {
 
 	config Config
 	state  *State
-	// initialised is closed when the machine is initialised as a member of a cluster.
+	// initialised is signalled when the machine is configured as a member of a cluster.
 	initialised chan struct{}
+	wgNetwork   *network.WireGuardNetwork
 
 	localServer   *grpc.Server
 	networkServer *grpc.Server
 
 	clusterState *cluster.State
-	cluster      *cluster.Server
+	cluster      *cluster.Cluster
 }
 
 func NewMachine(config *Config) (*Machine, error) {
@@ -77,7 +78,7 @@ func NewMachine(config *Config) (*Machine, error) {
 	m := &Machine{
 		config:      *config,
 		state:       state,
-		initialised: make(chan struct{}),
+		initialised: make(chan struct{}, 1),
 
 		localServer:   grpc.NewServer(),
 		networkServer: grpc.NewServer(),
@@ -85,21 +86,23 @@ func NewMachine(config *Config) (*Machine, error) {
 	pb.RegisterMachineServer(m.localServer, m)
 	pb.RegisterMachineServer(m.networkServer, m)
 
-	clusterStatePath := cluster.StatePath(config.DataDir)
+	clusterState := cluster.NewState(cluster.StatePath(config.DataDir))
-	clusterState := cluster.NewState(clusterStatePath)
 	if err = clusterState.Load(); err != nil {
-		if !errors.Is(err, os.ErrNotExist) {
+		if errors.Is(err, os.ErrNotExist) {
+			// Cluster state file does not exist, initialise the cluster without a state to fail cluster requests.
+			m.cluster = cluster.NewCluster(nil)
+		} else {
 			return nil, fmt.Errorf("load cluster state: %w", err)
 		}
 	} else {
-		// Cluster state is successfully loaded, start the cluster server.
+		// Cluster state is successfully loaded, initialise the cluster with it.
-		m.cluster = cluster.NewServer(clusterState)
+		m.cluster = cluster.NewCluster(clusterState)
+	}
 	pb.RegisterClusterServer(m.localServer, m.cluster)
 	pb.RegisterClusterServer(m.networkServer, m.cluster)
-	}
 
 	if m.IsInitialised() {
-		close(m.initialised)
+		m.initialised <- struct{}{}
 	}
 
 	return m, nil
@@ -134,8 +137,33 @@ func (m *Machine) Run(ctx context.Context) error {
 		},
 	)
 
-	// Start the machine network API server if the management IP is configured for it.
+	// Start the WireGuard network and network server after the machine is initialised as a member of a cluster.
-	if m.state.Network.ManagementIP != (netip.Addr{}) {
+	errGroup.Go(
+		func() error {
+			if !m.IsInitialised() {
+				slog.Info(
+					"Waiting for the machine to be initialised as a member of a cluster to start WireGuard network.",
+				)
+			}
+
+			netCancel := func() {}
+			for {
+				select {
+				case <-m.initialised:
+				case <-ctx.Done():
+					return nil
+				}
+
+				// Cancel the previously running network goroutine before reconfiguring the network.
+				netCancel()
+				wasConfigured := m.wgNetwork != nil
+				if err := m.configureNetwork(); err != nil {
+					return err
+				}
+
+				// Start the machine network API server if it was not already started.
+				// TODO: implement a proper mechanism to restart the network API server if the management IP changes.
+				if !wasConfigured {
 					apiAddr := net.JoinHostPort(m.state.Network.ManagementIP.String(), strconv.Itoa(APIPort))
 					networkListener, err := net.Listen("tcp", apiAddr)
 					if err != nil {
@@ -153,45 +181,23 @@ func (m *Machine) Run(ctx context.Context) error {
 					)
 				}
 
-	// Start the WireGuard network after the machine is initialised as a member of a cluster.
 				errGroup.Go(
 					func() error {
-			if !m.IsInitialised() {
+						var netCtx context.Context
-				slog.Info(
+						netCtx, netCancel = context.WithCancel(ctx)
-					"Waiting for the machine to be initialised as a member of a cluster to start WireGuard network.",
+						if err = m.wgNetwork.Run(netCtx); err != nil {
-				)
-			}
-			select {
-			case <-m.initialised:
-			case <-ctx.Done():
-				return nil
-			}
-
-			slog.Info("Starting WireGuard network.")
-			wgnet, err := network.NewWireGuardNetwork()
-			if err != nil {
-				return fmt.Errorf("create WireGuard network: %w", err)
-			}
-			if err = wgnet.Configure(*m.state.Network); err != nil {
-				return fmt.Errorf("configure WireGuard network: %w", err)
-			}
-
-			//ctx, cancel := context.WithCancel(context.Background())
-			//go wgnet.WatchEndpoints(ctx, peerEndpointChangeNotifier)
-
-			//addrs, err := network.ListRoutableIPs()
-			//if err != nil {
-			//	return err
-			//}
-			//fmt.Println("Addresses:", addrs)
-
-			if err = wgnet.Run(ctx); err != nil {
 							return fmt.Errorf("WireGuard network failed: %w", err)
 						}
 						return nil
 					},
 				)
 
+				//ctx, cancel := context.WithCancel(context.Background())
+				//go wgnet.WatchEndpoints(ctx, peerEndpointChangeNotifier)
+			}
+		},
+	)
+
 	// Shutdown goroutine.
 	errGroup.Go(
 		func() error {
@@ -212,6 +218,23 @@ func (m *Machine) Run(ctx context.Context) error {
 	return errGroup.Wait()
 }
 
+func (m *Machine) configureNetwork() error {
+	if m.wgNetwork == nil {
+		slog.Info("Starting WireGuard network.")
+		var err error
+		m.wgNetwork, err = network.NewWireGuardNetwork()
+		if err != nil {
+			return fmt.Errorf("create WireGuard network: %w", err)
+		}
+	}
+
+	if err := m.wgNetwork.Configure(*m.state.Network); err != nil {
+		return fmt.Errorf("configure WireGuard network: %w", err)
+	}
+	slog.Info("WireGuard network configured.")
+	return nil
+}
+
 // listenUnixSocket creates a new Unix socket listener with the specified path. The socket file is created with 0660
 // access mode and uncloud group if the group is found, otherwise it falls back to the root group.
 func listenUnixSocket(path string) (net.Listener, error) {
@@ -248,16 +271,15 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 	}
 
 	// TODO: a proper cluster leave mechanism and machine reset should be implemented later.
-	//  For now assume the cluster server is not running.
+	//  For now just reset the machine state and cluster state.
 	clusterStatePath := cluster.StatePath(m.config.DataDir)
 	clusterState := cluster.NewState(clusterStatePath)
 	if err = clusterState.Save(); err != nil {
 		return nil, status.Errorf(codes.Internal, "save cluster state: %v", err)
 	}
-	clusterServer := cluster.NewServer(clusterState)
+	m.cluster.SetState(clusterState)
-	// TODO: register and start the cluster server.
+	slog.Info("Cluster state initialised.", "path", clusterStatePath)
-
+	if err = m.cluster.SetNetwork(req.Network); err != nil {
-	if err = clusterServer.SetNetwork(req.Network); err != nil {
 		return nil, status.Errorf(codes.InvalidArgument, "set cluster network: %v", err)
 	}
 
@@ -286,7 +308,7 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 			PublicKey: m.state.Network.PublicKey,
 		},
 	}
-	addResp, err := clusterServer.AddMachine(ctx, addReq)
+	addResp, err := m.cluster.AddMachine(ctx, addReq)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "add machine to cluster: %v", err)
 	}
@@ -311,7 +333,7 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 
 	// Add a user to the cluster and build a peers config from it if provided.
 	if req.User != nil {
-		if err = clusterServer.AddUser(req.User); err != nil {
+		if err = m.cluster.AddUser(req.User); err != nil {
 			return nil, status.Errorf(codes.Internal, "add user to cluster: %v", err)
 		}
 		userManageIP, uErr := req.User.Network.ManagementIp.ToAddr()
@@ -329,9 +351,9 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 	if err = m.state.Save(); err != nil {
 		return nil, status.Errorf(codes.Internal, "save machine state: %v", err)
 	}
-	slog.Info("Cluster initialised.", "machine", m.state.Name)
+	slog.Info("Cluster initialised with machine.", "machine", m.state.Name)
 	// Signal that the machine is initialised as a member of a cluster.
-	close(m.initialised)
+	m.initialised <- struct{}{}
 
 	resp := &pb.InitClusterResponse{
 		Machine: addResp.Machine,