refactor network operations into networkController to start/stop when joining/leaving cluster

internal/machine/cluster/cluster.go

@@ -27,7 +27,7 @@ type Cluster struct {
 func NewCluster(state *State) *Cluster {
 	return &Cluster{
 		state:         state,
-		newMachinesCh: make(chan *pb.MachineInfo),
+		newMachinesCh: make(chan *pb.MachineInfo, 1),
 	}
 }
 
@@ -158,7 +158,6 @@ func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*p
 
 	// TODO: notify all cluster machines about the new machine so they can update their peers config.
 	//  In PoC we just notify the local machine.
-	// TODO: there is a race condition with network configuration if a new cluster is initialized on the machine.
 	c.newMachinesCh <- m
 
 	resp := &pb.AddMachineResponse{Machine: m}

internal/machine/machine.go

@@ -39,12 +39,8 @@ type Machine struct {
 	state  *State
 	// 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.Cluster
 	newMachinesCh <-chan *pb.MachineInfo
 }
@@ -77,32 +73,28 @@ func NewMachine(config *Config) (*Machine, error) {
 		}
 	}
 
-	m := &Machine{
-		config:      *config,
-		state:       state,
-		initialised: make(chan struct{}, 1),
-
-		localServer:   grpc.NewServer(),
-		networkServer: grpc.NewServer(),
-	}
-	pb.RegisterMachineServer(m.localServer, m)
-	pb.RegisterMachineServer(m.networkServer, m)
-
+	var c *cluster.Cluster
 	clusterState := cluster.NewState(cluster.StatePath(config.DataDir))
 	if err = clusterState.Load(); err != nil {
 		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)
+			c = cluster.NewCluster(nil)
 		} else {
 			return nil, fmt.Errorf("load cluster state: %w", err)
 		}
 	} else {
 		// Cluster state is successfully loaded, initialise the cluster with it.
-		m.cluster = cluster.NewCluster(clusterState)
+		c = cluster.NewCluster(clusterState)
 	}
-	pb.RegisterClusterServer(m.localServer, m.cluster)
-	pb.RegisterClusterServer(m.networkServer, m.cluster)
-	m.newMachinesCh = m.cluster.WatchNewMachines()
+
+	m := &Machine{
+		config:        *config,
+		state:         state,
+		initialised:   make(chan struct{}, 1),
+		cluster:       c,
+		newMachinesCh: c.WatchNewMachines(),
+	}
+	m.localServer = newGRPCServer(m, c)
 
 	if m.IsInitialised() {
 		m.initialised <- struct{}{}
@@ -111,9 +103,19 @@ func NewMachine(config *Config) (*Machine, error) {
 	return m, nil
 }
 
+func newGRPCServer(m pb.MachineServer, c pb.ClusterServer) *grpc.Server {
+	s := grpc.NewServer()
+	pb.RegisterMachineServer(s, m)
+	pb.RegisterClusterServer(s, c)
+	return s
+}
+
 // IsInitialised returns true if the machine has been configured as a member of a cluster,
 // either by initialising a new cluster on it or joining an existing one.
 func (m *Machine) IsInitialised() bool {
+	m.state.mu.RLock()
+	defer m.state.mu.RUnlock()
+
 	return m.state.ID != ""
 }
 
@@ -140,106 +142,53 @@ func (m *Machine) Run(ctx context.Context) error {
 		},
 	)
 
-	// Start the WireGuard network and network server after the machine is initialised as a member of a cluster.
+	// Control loop for managing the network controller.
 	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.",
+					"Waiting for the machine to be initialised as a member of a cluster " +
+						"to start the network controller.",
 				)
 			}
 
-			netCancel := func() {}
+			var ctrl *networkController
+			// Error channel for communicating the termination of the network controller.
+			errCh := make(chan error)
+
 			for {
 				select {
+				// Wait for the machine to be initialised as a member of a cluster to start the network controller.
+				// It can be reset when leaving the cluster and then re-initialised again with a new configuration.
 				case <-m.initialised:
-				case <-ctx.Done():
-					return nil
+					var err error
+					slog.Info("Starting network controller.")
+					networkServer := newGRPCServer(m, m.cluster)
+					ctrl, err = newNetworkController(m.state, networkServer, m.newMachinesCh)
+					if err != nil {
+						return fmt.Errorf("initialise network controller: %w", err)
 					}
 
-				// Cancel the previously running network goroutine before reconfiguring the network.
-				netCancel()
-				wasConfigured := m.wgNetwork != nil
-				if err := m.configureNetwork(); err != nil {
+					go func() {
+						if err = ctrl.Run(ctx); err != nil {
+							errCh <- fmt.Errorf("run network controller: %w", err)
+						} else {
+							slog.Info("Network controller stopped.")
+							errCh <- nil
+						}
+					}()
+				case err := <-errCh:
+					if 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 {
-						return fmt.Errorf("listen API port: %w", err)
-					}
-
-					errGroup.Go(
-						func() error {
-							slog.Info("Starting network API server.", "addr", apiAddr)
-							if err = m.networkServer.Serve(networkListener); err != nil {
-								return fmt.Errorf("network API server failed: %w", err)
-							}
-							return nil
-						},
-					)
-				}
-
-				errGroup.Go(
-					func() error {
-						var netCtx context.Context
-						netCtx, netCancel = context.WithCancel(ctx)
-						if err = m.wgNetwork.Run(netCtx); err != nil {
-							return fmt.Errorf("WireGuard network failed: %w", err)
-						}
-						return nil
-					},
-				)
-
-				//ctx, cancel := context.WithCancel(context.Background())
-				//go wgnet.WatchEndpoints(ctx, peerEndpointChangeNotifier)
-			}
-		},
-	)
-
-	// Handle new machines added to the cluster.
-	errGroup.Go(
-		func() error {
-			for {
-				select {
-				case machineInfo := <-m.newMachinesCh:
-					slog.Info("Handling new machine added to the cluster.", "name", machineInfo.Name)
-					if err := machineInfo.Network.Validate(); err != nil {
-						slog.Error("Invalid machine network configuration.", "err", err)
-						continue
-					}
-					// Ignore errors as they are already validated.
-					subnet, _ := machineInfo.Network.Subnet.ToPrefix()
-					manageIP, _ := machineInfo.Network.ManagementIp.ToAddr()
-					endpoints := make([]netip.AddrPort, len(machineInfo.Network.Endpoints))
-					for i, ep := range machineInfo.Network.Endpoints {
-						addrPort, _ := ep.ToAddrPort()
-						endpoints[i] = addrPort
-					}
-
-					peer := network.PeerConfig{
-						Subnet:       &subnet,
-						ManagementIP: manageIP,
-						AllEndpoints: endpoints,
-						PublicKey:    machineInfo.Network.PublicKey,
-					}
-					if len(endpoints) > 0 {
-						peer.Endpoint = &endpoints[0]
-					}
-
-					m.state.Network.Peers = append(m.state.Network.Peers, peer)
-					if err := m.state.Save(); err != nil {
-						return fmt.Errorf("save machine state: %w", err)
-					}
-
-					if err := m.configureNetwork(); err != nil {
-						return fmt.Errorf("configure network with new peer: %w", err)
-					}
+					ctrl = nil
 				case <-ctx.Done():
+					// Wait for the network controller to stop before returning.
+					if ctrl != nil {
+						if err := <-errCh; err != nil {
+							return err
+						}
+					}
 					return nil
 				}
 			}
@@ -250,15 +199,10 @@ func (m *Machine) Run(ctx context.Context) error {
 	errGroup.Go(
 		func() error {
 			<-ctx.Done()
-			slog.Info("Stopping network API server.")
-			// TODO: implement timeout for graceful shutdown.
-			m.networkServer.GracefulStop()
-			slog.Info("network API server stopped.")
-
 			slog.Info("Stopping local API server.")
 			// TODO: implement timeout for graceful shutdown.
 			m.localServer.GracefulStop()
-			slog.Info("local API server stopped.")
+			slog.Info("Local API server stopped.")
 			return nil
 		},
 	)
@@ -266,23 +210,6 @@ 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) {
@@ -402,8 +329,6 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 	slog.Info("Cluster initialised with machine.", "machine", m.state.Name)
 	// Signal that the machine is initialised as a member of a cluster.
 	m.initialised <- struct{}{}
-	// TODO: consider calling a synchronous method to reconfigure the network to return error if it fails.
-	//  Alternatively a client can call another method to check the network status.
 
 	resp := &pb.InitClusterResponse{
 		Machine: addResp.Machine,
@@ -477,8 +402,6 @@ func (m *Machine) JoinCluster(ctx context.Context, req *pb.JoinClusterRequest) (
 	slog.Info("Machine configured to join the cluster.", "id", m.state.ID, "name", m.state.Name)
 	// Signal that the machine is initialised as a member of a cluster.
 	m.initialised <- struct{}{}
-	// TODO: consider calling a synchronous method to reconfigure the network to return error if it fails.
-	//  Alternatively a client can call another method to check the network status.
 
 	return &emptypb.Empty{}, nil
 }

internal/machine/network.go

@@ -0,0 +1,160 @@
+package machine
+
+import (
+	"context"
+	"fmt"
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"log/slog"
+	"net"
+	"net/netip"
+	"strconv"
+	"uncloud/internal/machine/api/pb"
+	"uncloud/internal/machine/network"
+)
+
+const APIPort = 51000
+
+type networkController struct {
+	state         *State
+	wgnet         *network.WireGuardNetwork
+	server        *grpc.Server
+	newMachinesCh <-chan *pb.MachineInfo
+	// TODO: DNS server/resolver
+}
+
+func newNetworkController(state *State, server *grpc.Server, newMachCh <-chan *pb.MachineInfo) (
+	*networkController, error,
+) {
+	slog.Info("Starting WireGuard network.")
+	wgnet, err := network.NewWireGuardNetwork()
+	if err != nil {
+		return nil, fmt.Errorf("create WireGuard network: %w", err)
+	}
+
+	return &networkController{
+		state:         state,
+		wgnet:         wgnet,
+		server:        server,
+		newMachinesCh: newMachCh,
+	}, nil
+}
+
+func (nc *networkController) Run(ctx context.Context) error {
+	if err := nc.wgnet.Configure(*nc.state.Network); err != nil {
+		return fmt.Errorf("configure WireGuard network: %w", err)
+	}
+	slog.Info("WireGuard network configured.")
+
+	errGroup, ctx := errgroup.WithContext(ctx)
+
+	// Start the network API server. Assume the management IP can't be changed when the network is running.
+	apiAddr := net.JoinHostPort(nc.state.Network.ManagementIP.String(), strconv.Itoa(APIPort))
+	listener, err := net.Listen("tcp", apiAddr)
+	if err != nil {
+		return fmt.Errorf("listen API port: %w", err)
+	}
+	errGroup.Go(
+		func() error {
+			slog.Info("Starting network API server.", "addr", apiAddr)
+			if err := nc.server.Serve(listener); err != nil {
+				return fmt.Errorf("network API server failed: %w", err)
+			}
+			return nil
+		},
+	)
+
+	// Handle new machines added to the cluster. Handling new machines and endpoint changes should be done
+	// in separate goroutines to avoid a deadlock when reconfiguring the network.
+	errGroup.Go(
+		func() error {
+			if err := nc.handleNewMachines(ctx); err != nil {
+				return fmt.Errorf("handle new machines: %w", err)
+			}
+			return nil
+		},
+	)
+
+	// TODO: run another goroutine to watch WG network endpoint changes and update the state accordingly.
+	//  Network updates in the state should not occur outside of this controller.
+
+	errGroup.Go(
+		func() error {
+			if err = nc.wgnet.Run(ctx); err != nil {
+				return fmt.Errorf("WireGuard network failed: %w", err)
+			}
+			return nil
+		},
+	)
+
+	// Wait for the context to be done and stop the network API server.
+	errGroup.Go(
+		func() error {
+			<-ctx.Done()
+			slog.Info("Stopping network API server.")
+			// TODO: implement timeout for graceful shutdown.
+			nc.server.GracefulStop()
+			slog.Info("Network API server stopped.")
+			return nil
+		},
+	)
+
+	return errGroup.Wait()
+}
+
+func (nc *networkController) handleNewMachines(ctx context.Context) error {
+	for {
+		select {
+		case minfo := <-nc.newMachinesCh:
+			slog.Info("Handling new machine added to the cluster.", "name", minfo.Name)
+
+			// Skip the current machine.
+			nc.state.mu.RLock()
+			currentMachID := nc.state.ID
+			nc.state.mu.RUnlock()
+			if minfo.Id == currentMachID {
+				continue
+			}
+
+			if err := minfo.Network.Validate(); err != nil {
+				slog.Error("Invalid machine network configuration.", "err", err)
+				continue
+			}
+			// Ignore errors as they are already validated.
+			subnet, _ := minfo.Network.Subnet.ToPrefix()
+			manageIP, _ := minfo.Network.ManagementIp.ToAddr()
+			endpoints := make([]netip.AddrPort, len(minfo.Network.Endpoints))
+			for i, ep := range minfo.Network.Endpoints {
+				addrPort, _ := ep.ToAddrPort()
+				endpoints[i] = addrPort
+			}
+
+			peer := network.PeerConfig{
+				Subnet:       &subnet,
+				ManagementIP: manageIP,
+				AllEndpoints: endpoints,
+				PublicKey:    minfo.Network.PublicKey,
+			}
+			if len(endpoints) > 0 {
+				peer.Endpoint = &endpoints[0]
+			}
+
+			nc.state.mu.Lock()
+			// TODO: deduplicate peers by public key, maybe implement addNetworkPeer method in the state.
+			nc.state.Network.Peers = append(nc.state.Network.Peers, peer)
+			err := nc.state.Save()
+			nc.state.mu.Unlock()
+			if err != nil {
+				return fmt.Errorf("save machine state: %w", err)
+			}
+
+			if err := nc.wgnet.Configure(*nc.state.Network); err != nil {
+				return fmt.Errorf("configure network with new peer: %w", err)
+			}
+		case <-ctx.Done():
+			return nil
+		}
+	}
+}
+
+// TODO: method to shutdown network when leaving a cluster. Regular context cancellation shouldn't bring it down.

internal/machine/network/wireguard_linux.go

@@ -13,12 +13,15 @@ import (
 	"log/slog"
 	"net/netip"
 	"slices"
+	"sync"
 	"time"
 )
 
 type WireGuardNetwork struct {
 	link  netlink.Link
 	peers []peer
+	// mu synchronises concurrent network configuration changes.
+	mu sync.Mutex
 }
 
 type peer struct {
@@ -68,6 +71,9 @@ func createOrGetLink(name string) (netlink.Link, error) {
 // Configure applies the given configuration to the WireGuard network interface.
 // It updates device and peers settings, subnet, and peer routes.
 func (n *WireGuardNetwork) Configure(config Config) error {
+	n.mu.Lock()
+	defer n.mu.Unlock()
+
 	// Reconstruct the list of peers, ensuring that the last endpoint change time is preserved for any existing peers.
 	existingPeersByPublicKey := map[string]peer{}
 	for _, p := range n.peers {
@@ -112,8 +118,10 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 	if err = n.updateAddresses(addrs); err != nil {
 		return err
 	}
-	slog.Info("Updated addresses of the WireGuard interface.",
-		"name", n.link.Attrs().Name, "addrs", addrs)
+	slog.Info(
+		"Updated addresses of the WireGuard interface.",
+		"name", n.link.Attrs().Name, "addrs", addrs,
+	)
 
 	// Bring the WireGuard interface up if it's not already up.
 	if n.link.Attrs().Flags&unix.IFF_UP != unix.IFF_UP {
@@ -125,8 +133,10 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 	if err = n.updatePeerRoutes(); err != nil {
 		return err
 	}
-	slog.Info("Updated routes to peers via the WireGuard interface.",
-		"name", n.link.Attrs().Name, "peers", len(n.peers))
+	slog.Info(
+		"Updated routes to peers via the WireGuard interface.",
+		"name", n.link.Attrs().Name, "peers", len(n.peers),
+	)
 
 	return nil
 }
@@ -148,9 +158,11 @@ func (n *WireGuardNetwork) updateAddresses(addrs []netip.Prefix) error {
 		return fmt.Errorf("list addresses on WireGuard link %q: %w", n.link.Attrs().Name, err)
 	}
 	for _, linkAddr := range linkAddrs {
-		if slices.ContainsFunc(addrs, func(a netip.Prefix) bool {
+		if slices.ContainsFunc(
+			addrs, func(a netip.Prefix) bool {
 				return linkAddr.IPNet.String() == a.String()
-		}) {
+			},
+		) {
 			continue
 		}
 		if err = netlink.AddrDel(n.link, &linkAddr); err != nil {
@@ -182,15 +194,19 @@ func (n *WireGuardNetwork) updatePeerRoutes() error {
 	// Add routes to the computed IP ranges via the WireGuard link.
 	for _, prefix := range ipset.Prefixes() {
 		dst := prefixToIPNet(prefix)
-		if err = netlink.RouteAdd(&netlink.Route{
+		if err = netlink.RouteAdd(
+			&netlink.Route{
 				LinkIndex: n.link.Attrs().Index,
 				Scope:     netlink.SCOPE_LINK,
 				Dst:       &dst,
-		}); err != nil && !errors.Is(err, unix.EEXIST) {
+			},
+		); err != nil && !errors.Is(err, unix.EEXIST) {
 			return fmt.Errorf("add route to WireGuard link %q: %w", n.link.Attrs().Name, err)
 		}
-		slog.Debug("Added route to peer(s) via WireGuard interface.",
-			"name", n.link.Attrs().Name, "dst", prefix)
+		slog.Debug(
+			"Added route to peer(s) via WireGuard interface.",
+			"name", n.link.Attrs().Name, "dst", prefix,
+		)
 	}
 
 	// Remove old routes to IP ranges that are no longer in the configuration.
@@ -210,8 +226,10 @@ func (n *WireGuardNetwork) updatePeerRoutes() error {
 		if err = netlink.RouteDel(&route); err != nil {
 			return fmt.Errorf("remove route %q from WireGuard link %q: %w", route.Dst, n.link.Attrs().Name, err)
 		}
-		slog.Debug("Removed route to peer(s) via WireGuard interface.",
-			"name", n.link.Attrs().Name, "dst", routePrefix)
+		slog.Debug(
+			"Removed route to peer(s) via WireGuard interface.",
+			"name", n.link.Attrs().Name, "dst", routePrefix,
+		)
 	}
 	return nil
 }

internal/machine/state.go

@@ -6,6 +6,7 @@ import (
 	"net/netip"
 	"os"
 	"path/filepath"
+	"sync"
 	"uncloud/internal/machine/cluster"
 	"uncloud/internal/machine/network"
 )
@@ -13,7 +14,6 @@ import (
 const (
 	DefaultDataDir = "/var/lib/uncloud"
 	StateFileName  = "machine.json"
-	APIPort        = 51000
 )
 
 // State defines the machine-specific configuration within a cluster. It encapsulates essential identifiers
@@ -28,6 +28,8 @@ type State struct {
 
 	// path is the file path config is read from and saved to.
 	path string
+	// mu protects the state from concurrent reads and writes.
+	mu sync.RWMutex
 }
 
 // StatePath returns the path to the machine state file within the given data directory.