refactor daemon to use Machine to manage all components

cmd/uncloud/machine/token.go

@@ -3,8 +3,8 @@ package machine
 import (
 	"fmt"
 	"github.com/spf13/cobra"
+	"uncloud/internal/daemon"
 	"uncloud/internal/machine"
-	"uncloud/internal/machine/daemon"
 )
 
 type tokenOptions struct {

internal/daemon/daemon.go

@@ -2,14 +2,9 @@ package daemon
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"golang.org/x/sync/errgroup"
 	"log/slog"
-	"net"
 	"net/netip"
-	"os"
-	"strconv"
 	"uncloud/internal/machine"
 	"uncloud/internal/machine/api/pb"
 	"uncloud/internal/machine/cluster"
@@ -108,119 +103,23 @@ func InitCluster(dataDir, machineName string, netPrefix netip.Prefix, users []*p
 
 type Daemon struct {
 	machine *machine.Machine
-
-	state   *machine.State
-	cluster *cluster.Server
 }
 
 func New(dataDir string) (*Daemon, error) {
-	mstatePath := machine.StatePath(dataDir)
-	mstate, err := machine.ParseState(mstatePath)
+	config := &machine.Config{
+		DataDir: dataDir,
+	}
+	mach, err := machine.NewMachine(config)
 	if err != nil {
-		if !errors.Is(err, os.ErrNotExist) {
-			return nil, fmt.Errorf("load machine config: %w", err)
-		}
-		// Generate an empty machine config with a new key pair.
-		slog.Info("Machine config not found, creating a new one.", "path", mstatePath)
-		privKey, pubKey, kErr := network.NewMachineKeys()
-		if kErr != nil {
-			return nil, fmt.Errorf("generate machine keys: %w", kErr)
-		}
-		slog.Info("Generated machine key pair.", "pubkey", pubKey)
-
-		mstate = &machine.State{
-			Network: &network.Config{
-				PrivateKey: privKey,
-				PublicKey:  pubKey,
-			},
-		}
-		mstate.SetPath(mstatePath)
-		if err = mstate.Save(); err != nil {
-			return nil, fmt.Errorf("save machine config: %w", err)
-		}
+		return nil, fmt.Errorf("init machine: %w", err)
 	}
 
-	cstatePath := cluster.StatePath(dataDir)
-	cstate := cluster.NewState(cstatePath)
-	if err = cstate.Load(); err != nil {
-		if !errors.Is(err, os.ErrNotExist) {
-			return nil, fmt.Errorf("load cluster state: %w", err)
-		}
-		slog.Info("Cluster state not found, creating a new one.", "path", cstatePath)
-		if err = cstate.Save(); err != nil {
-			return nil, fmt.Errorf("save cluster state: %w", err)
-		}
-	}
-
-	d := &Daemon{
-		state: mstate,
-	}
-	if mstate.Network.IsConfigured() {
-		config := &machine.Config{
-			APIAddr: net.JoinHostPort(mstate.Network.ManagementIP.String(), strconv.Itoa(machine.APIPort)),
-		}
-		d.machine, err = machine.NewMachine(config)
-		if err != nil {
-			return nil, fmt.Errorf("init machine: %w", err)
-		}
-	}
-
-	return d, nil
+	return &Daemon{
+		machine: mach,
+	}, nil
 }
 
 func (d *Daemon) Run(ctx context.Context) error {
-	// Use an errgroup to coordinate error handling and graceful shutdown of multiple daemon components.
-	errGroup, ctx := errgroup.WithContext(ctx)
-
-	// Start the network only if it is configured.
-	if d.state.Network.IsConfigured() {
-		wgnet, err := network.NewWireGuardNetwork()
-		if err != nil {
-			return fmt.Errorf("create WireGuard network: %w", err)
-		}
-		if err = wgnet.Configure(*d.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)
-
-		errGroup.Go(func() error {
-			if err = wgnet.Run(ctx); err != nil {
-				return fmt.Errorf("WireGuard network failed: %w", err)
-			}
-			return nil
-		})
-	} else {
-		slog.Info("Waiting for network configuration to start WireGuard network.")
-	}
-
-	if d.cluster != nil {
-		errGroup.Go(func() error {
-			slog.Info("Starting cluster.")
-			if err := d.machine.Run(); err != nil {
-				return fmt.Errorf("cluster failed: %w", err)
-			}
-			return nil
-		})
-	}
-
-	// Shutdown goroutine.
-	errGroup.Go(func() error {
-		<-ctx.Done()
-		if d.cluster != nil {
-			slog.Info("Stopping cluster.")
-			d.machine.Stop()
-			slog.Info("Cluster server stopped.")
-		}
-		return nil
-	})
-
-	return errGroup.Wait()
+	slog.Info("Starting machine.")
+	return d.machine.Run(ctx)
 }

internal/machine/machine.go

@@ -1,54 +1,147 @@
 package machine
 
 import (
+	"context"
+	"errors"
 	"fmt"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"log/slog"
 	"net"
+	"net/netip"
+	"os"
+	"strconv"
 	"uncloud/internal/machine/api/pb"
 	"uncloud/internal/machine/cluster"
+	"uncloud/internal/machine/network"
 )
 
 type Config struct {
 	// DataDir is the directory where the machine stores its persistent state.
 	DataDir     string
-	APIAddr     string
 	APISockPath string
 }
 
 type Machine struct {
 	config Config
-	server *grpc.Server
+	state  *State
+
+	networkServer *grpc.Server
+	cluster       *cluster.Server
+	// TODO: create localServer for unix socket.
 }
 
 func NewMachine(config *Config) (*Machine, error) {
-	m := &Machine{
-		config: *config,
-		server: grpc.NewServer(),
-	}
-
-	clusterState := cluster.NewState(cluster.StatePath(config.DataDir))
-	clusterServer := cluster.NewServer(clusterState)
-	pb.RegisterClusterServer(m.server, clusterServer)
-
-	return m, nil
-}
-
-func (m *Machine) Run() error {
-	listener, err := net.Listen("tcp", m.config.APIAddr)
+	// Load the existing machine state or create a new one.
+	statePath := StatePath(config.DataDir)
+	state, err := ParseState(statePath)
 	if err != nil {
-		return fmt.Errorf("listen API port: %w", err)
+		if !errors.Is(err, os.ErrNotExist) {
+			return nil, fmt.Errorf("load machine state: %w", err)
+		}
+		// Generate an empty machine config with a new key pair.
+		slog.Info("Machine state file not found, creating a new one.", "path", statePath)
+		privKey, pubKey, kErr := network.NewMachineKeys()
+		if kErr != nil {
+			return nil, fmt.Errorf("generate machine keys: %w", kErr)
+		}
+		slog.Info("Generated machine key pair.", "pubkey", pubKey)
+
+		state = &State{
+			Network: &network.Config{
+				PrivateKey: privKey,
+				PublicKey:  pubKey,
+			},
+		}
+		state.SetPath(statePath)
+		if err = state.Save(); err != nil {
+			return nil, fmt.Errorf("save machine state: %w", err)
+		}
 	}
-	slog.Info("Starting API server.", "addr", m.config.APIAddr)
-	if err = m.server.Serve(listener); err != nil {
-		return fmt.Errorf("API server failed: %w", err)
+
+	clusterStatePath := cluster.StatePath(config.DataDir)
+	clusterState := cluster.NewState(clusterStatePath)
+	if err = clusterState.Load(); err != nil {
+		if !errors.Is(err, os.ErrNotExist) {
+			return nil, fmt.Errorf("load cluster state: %w", err)
+		}
+		slog.Info("Cluster state file not found, creating a new one.", "path", clusterStatePath)
+		if err = clusterState.Save(); err != nil {
+			return nil, fmt.Errorf("save cluster state: %w", err)
+		}
 	}
-	return nil
+	clusterServer := cluster.NewServer(clusterState)
+
+	networkServer := grpc.NewServer()
+	pb.RegisterClusterServer(networkServer, clusterServer)
+
+	return &Machine{
+		config:        *config,
+		state:         state,
+		networkServer: networkServer,
+		cluster:       clusterServer,
+	}, nil
 }
 
-func (m *Machine) Stop() {
-	slog.Info("Stopping API server.")
-	// TODO: implement timeout for graceful shutdown.
-	m.server.GracefulStop()
-	slog.Info("API server stopped.")
+func (m *Machine) Run(ctx context.Context) error {
+	// Use an errgroup to coordinate error handling and graceful shutdown of multiple machine components.
+	errGroup, ctx := errgroup.WithContext(ctx)
+
+	// Start the network only if it is configured.
+	if m.state.Network.IsConfigured() {
+		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)
+
+		errGroup.Go(func() error {
+			if err = wgnet.Run(ctx); err != nil {
+				return fmt.Errorf("WireGuard network failed: %w", err)
+			}
+			return nil
+		})
+	} else {
+		slog.Info("Waiting for network configuration to start WireGuard network.")
+	}
+
+	// Start the machine API server if the management IP is configured for it.
+	if m.state.Network.ManagementIP != (netip.Addr{}) {
+		apiAddr := net.JoinHostPort(m.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 API server.", "addr", apiAddr)
+			if err = m.networkServer.Serve(listener); err != nil {
+				return fmt.Errorf("API server failed: %w", err)
+			}
+			return nil
+		})
+	}
+
+	// Shutdown goroutine.
+	errGroup.Go(func() error {
+		<-ctx.Done()
+		slog.Info("Stopping API server.")
+		// TODO: implement timeout for graceful shutdown.
+		m.networkServer.GracefulStop()
+		slog.Info("API server stopped.")
+		return nil
+	})
+
+	return errGroup.Wait()
 }