update network controller to reconfigure network peers on machine changes

internal/machine/cluster/cluster.go

@@ -21,15 +21,10 @@ type Cluster struct {
 	pb.UnimplementedClusterServer
 
 	store *store.Store
-	// TODO: temporary channel until the state is replaced with networkDB.
-	newMachinesCh chan *pb.MachineInfo
 }
 
 func NewCluster(store *store.Store) *Cluster {
-	return &Cluster{
+	return &Cluster{store: store}
-		store:         store,
-		newMachinesCh: make(chan *pb.MachineInfo, 1),
-	}
 }
 
 func (c *Cluster) Init(ctx context.Context, network netip.Prefix) error {
@@ -88,11 +83,6 @@ func (c *Cluster) Network(ctx context.Context) (netip.Prefix, error) {
 	return prefix, nil
 }
 
-// TODO: this is a temporary watcher for PoC until the state is state is replaced with networkDB.
-func (c *Cluster) WatchNewMachines() <-chan *pb.MachineInfo {
-	return c.newMachinesCh
-}
-
 // 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 {
@@ -175,10 +165,6 @@ func (c *Cluster) AddMachine(ctx context.Context, req *pb.AddMachineRequest) (*p
 	slog.Info("Machine added to the cluster.",
 		"id", m.Id, "name", m.Name, "subnet", subnet, "public_key", secret.Secret(m.Network.PublicKey))
 
-	// TODO: Subscribe all cluster members to updates about the new machine so they can update their peers config.
-	//  In PoC we just notify the local machine.
-	c.newMachinesCh <- m
-
 	resp := &pb.AddMachineResponse{Machine: m}
 	return resp, nil
 }

internal/machine/machine.go

@@ -79,9 +79,10 @@ type Machine struct {
 	// initialised is signalled when the machine is configured as a member of a cluster.
 	initialised chan struct{}
 
-	localServer   *grpc.Server
+	// store is the cluster store backed by a distributed Corrosion database.
-	cluster       *cluster.Cluster
+	store       *store.Store
-	newMachinesCh <-chan *pb.MachineInfo
+	cluster     *cluster.Cluster
+	localServer *grpc.Server
 }
 
 func NewMachine(config *Config) (*Machine, error) {
@@ -122,12 +123,12 @@ func NewMachine(config *Config) (*Machine, error) {
 	c := cluster.NewCluster(corroStore)
 
 	m := &Machine{
-		config:        *config,
+		config:      *config,
-		state:         state,
+		state:       state,
-		started:       make(chan struct{}),
+		started:     make(chan struct{}),
-		initialised:   make(chan struct{}, 1),
+		initialised: make(chan struct{}, 1),
-		cluster:       c,
+		store:       corroStore,
-		newMachinesCh: c.WatchNewMachines(),
+		cluster:     c,
 	}
 	m.localServer = newGRPCServer(m, c)
 
@@ -223,7 +224,7 @@ func (m *Machine) Run(ctx context.Context) error {
 
 					slog.Info("Starting network controller.")
 					networkServer := newGRPCServer(m, m.cluster)
-					ctrl, err = newNetworkController(m.state, networkServer, m.config.CorrosionService, m.newMachinesCh)
+					ctrl, err = newNetworkController(m.state, m.store, networkServer, m.config.CorrosionService)
 					if err != nil {
 						return fmt.Errorf("initialise network controller: %w", err)
 					}
@@ -430,7 +431,7 @@ func (m *Machine) InitCluster(ctx context.Context, req *pb.InitClusterRequest) (
 }
 
 // JoinCluster configures the local machine to join an existing cluster.
-func (m *Machine) JoinCluster(ctx context.Context, req *pb.JoinClusterRequest) (*emptypb.Empty, error) {
+func (m *Machine) JoinCluster(_ context.Context, req *pb.JoinClusterRequest) (*emptypb.Empty, error) {
 	if m.Initialised() {
 		return nil, status.Error(codes.FailedPrecondition, "machine is already configured as a cluster member")
 	}

internal/machine/network.go

@@ -2,16 +2,21 @@ package machine
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"github.com/cenkalti/backoff/v4"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"log/slog"
 	"net"
 	"net/netip"
+	"slices"
 	"strconv"
+	"time"
 	"uncloud/internal/machine/api/pb"
 	"uncloud/internal/machine/corroservice"
 	"uncloud/internal/machine/network"
+	"uncloud/internal/machine/store"
 )
 
 const (
@@ -21,17 +26,18 @@ const (
 )
 
 type networkController struct {
-	state         *State
+	state        *State
-	wgnet         *network.WireGuardNetwork
+	store        *store.Store
-	server        *grpc.Server
+	wgnet        *network.WireGuardNetwork
-	corroService  corroservice.Service
+	server       *grpc.Server
-	newMachinesCh <-chan *pb.MachineInfo
+	corroService corroservice.Service
+
 	// TODO: DNS server/resolver listening on the machine IP, e.g. 10.210.0.1:53. It can't listen on 127.0.X.X
 	//  like resolved does because it needs to be reachable from both the host and the containers.
 }
 
 func newNetworkController(
-	state *State, server *grpc.Server, corroService corroservice.Service, newMachCh <-chan *pb.MachineInfo,
+	state *State, store *store.Store, server *grpc.Server, corroService corroservice.Service,
 ) (
 	*networkController, error,
 ) {
@@ -42,11 +48,11 @@ func newNetworkController(
 	}
 
 	return &networkController{
-		state:         state,
+		state:        state,
-		wgnet:         wgnet,
+		store:        store,
-		server:        server,
+		wgnet:        wgnet,
-		corroService:  corroService,
+		server:       server,
-		newMachinesCh: newMachCh,
+		corroService: corroService,
 	}, nil
 }
 
@@ -98,11 +104,11 @@ func (nc *networkController) Run(ctx context.Context) error {
 		},
 	)
 
-	// Handle new machines added to the cluster. Handling new machines and endpoint changes should be done
+	// Handle machine changes in the cluster. Handling machine 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 {
+			if err := nc.handleMachineChanges(ctx); err != nil {
 				return fmt.Errorf("handle new machines: %w", err)
 			}
 			return nil
@@ -136,59 +142,117 @@ func (nc *networkController) Run(ctx context.Context) error {
 	return errGroup.Wait()
 }
 
-func (nc *networkController) handleNewMachines(ctx context.Context) error {
+// handleMachineChanges subscribes to machine changes in the cluster and reconfigures the network peers accordingly.
+func (nc *networkController) handleMachineChanges(ctx context.Context) error {
 	for {
-		select {
+		// Retry to subscribe to machine changes indefinitely until the context is done.
-		case minfo := <-nc.newMachinesCh:
+		b := backoff.WithContext(backoff.NewExponentialBackOff(
-			slog.Info("Handling new machine added to the cluster.", "name", minfo.Name)
+			backoff.WithInitialInterval(1*time.Second),
+			backoff.WithMaxInterval(60*time.Second),
+			backoff.WithMaxElapsedTime(0),
+		), ctx)
 
-			// Skip the current machine.
+		var (
-			nc.state.mu.RLock()
+			machines []*pb.MachineInfo
-			currentMachID := nc.state.ID
+			changes  <-chan struct{}
-			nc.state.mu.RUnlock()
+			err      error
-			if minfo.Id == currentMachID {
+		)
-				continue
+		subscribe := func() error {
+			if machines, changes, err = nc.store.SubscribeMachines(ctx); err != nil {
+				slog.Info("Failed to subscribe to machine changes, retrying.", "err", err)
 			}
+			return err
+		}
+		if err = backoff.Retry(subscribe, b); err != nil {
+			if errors.Is(err, context.Canceled) {
+				return nil
+			}
+			slog.Error("Unexpected error while retrying to subscribe to machine changes.", "err", err)
+			continue
+		}
+		slog.Info("Subscribed to machine changes in the cluster to reconfigure network peers.")
 
-			if err := minfo.Network.Validate(); err != nil {
+		if err = nc.configurePeers(machines); err != nil {
-				slog.Error("Invalid machine network configuration.", "err", err)
+			slog.Error("Failed to configure peers.", "err", err)
-				continue
+		}
+		// For simplicity, reconfigure all peers on any change.
+		for {
+			select {
+			case <-changes:
+				slog.Info("Cluster machines changed, reconfiguring network peers.")
+				if machines, err = nc.store.ListMachines(ctx); err != nil {
+					slog.Error("Failed to list machines.", "err", err)
+					continue
+				}
+				if err = nc.configurePeers(machines); err != nil {
+					slog.Error("Failed to configure peers.", "err", err)
+				}
+			case <-ctx.Done():
+				return nil
 			}
-			// 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
 		}
 	}
 }
 
+func (nc *networkController) configurePeers(machines []*pb.MachineInfo) error {
+	nc.state.mu.RLock()
+	currentPeerEndpoints := make(map[string]*netip.AddrPort, len(nc.state.Network.Peers))
+	for _, p := range nc.state.Network.Peers {
+		currentPeerEndpoints[p.PublicKey.String()] = p.Endpoint
+	}
+	nc.state.mu.RUnlock()
+
+	// Construct the list of peers from the machine configurations ensuring that the current endpoint is preserved.
+	peers := make([]network.PeerConfig, 0, len(machines)-1)
+	for _, m := range machines {
+		// Skip the current machine.
+		if m.Id == nc.state.ID {
+			continue
+		}
+		if err := m.Network.Validate(); err != nil {
+			slog.Error("Invalid machine network configuration.", "machine", m.Name, "err", err)
+			continue
+		}
+		// Ignore errors as they are already validated.
+		subnet, _ := m.Network.Subnet.ToPrefix()
+		manageIP, _ := m.Network.ManagementIp.ToAddr()
+		endpoints := make([]netip.AddrPort, len(m.Network.Endpoints))
+		for i, ep := range m.Network.Endpoints {
+			addrPort, _ := ep.ToAddrPort()
+			endpoints[i] = addrPort
+		}
+		peer := network.PeerConfig{
+			Subnet:       &subnet,
+			ManagementIP: manageIP,
+			AllEndpoints: endpoints,
+			PublicKey:    m.Network.PublicKey,
+		}
+
+		currentEndpoint := currentPeerEndpoints[peer.PublicKey.String()]
+		if currentEndpoint != nil && slices.Contains(endpoints, *currentEndpoint) {
+			peer.Endpoint = currentEndpoint
+		} else if len(endpoints) > 0 {
+			peer.Endpoint = &endpoints[0]
+		}
+
+		peers = append(peers, peer)
+	}
+
+	// Preserve the new list of peers in the machine state.
+	nc.state.mu.Lock()
+	nc.state.Network.Peers = peers
+	err := nc.state.Save()
+	nc.state.mu.Unlock()
+	if err != nil {
+		return fmt.Errorf("save machine state: %w", err)
+	}
+
+	nc.state.mu.RLock()
+	defer nc.state.mu.RUnlock()
+	if err = nc.wgnet.Configure(*nc.state.Network); err != nil {
+		return fmt.Errorf("configure network peers: %w", err)
+	}
+	return nil
+}
+
 // TODO: method to shutdown network when leaving a cluster. Regular context cancellation shouldn't bring it down.