chore: simplify cluster controller initialisation, prepare for reset

internal/machine/cluster.go

@@ -29,39 +29,38 @@ const (
 	APIPort = 51000
 )
 
-type networkController struct {
+// clusterController is the main controller for the machine that is a cluster member. It manages components such as
+// the WireGuard network, Corrosion service, Docker network and containers, and embedded DNS server.
+type clusterController struct {
 	state *State
 	store *store.Store
 
 	wgnet           *network.WireGuardNetwork
 	endpointChanges <-chan network.EndpointChangeEvent
 
-	server        *grpc.Server
-	corroService  corroservice.Service
-	dockerCli     *client.Client
-	caddyfileCtrl *caddyconfig.Controller
+	server       *grpc.Server
+	corroService corroservice.Service
+	dockerCli    *client.Client
+	// dockerReady is signalled when Docker is configured and ready for containers.
+	dockerReady     chan<- struct{}
+	caddyconfigCtrl *caddyconfig.Controller
 
 	// dnsServer is the embedded internal DNS server for the cluster listening on the machine IP.
 	dnsServer   *dns.Server
 	dnsResolver *dns.ClusterResolver
-
-	// networkReady is signalled when the Docker network is configured and ready for containers.
-	networkReady chan<- struct{}
 }
 
-func newNetworkController(
+func newClusterController(
 	state *State,
 	store *store.Store,
 	server *grpc.Server,
 	corroService corroservice.Service,
 	dockerCli *client.Client,
+	dockerReady chan<- struct{},
 	caddyfileCtrl *caddyconfig.Controller,
 	dnsServer *dns.Server,
 	dnsResolver *dns.ClusterResolver,
-	networkReady chan<- struct{},
-) (
-	*networkController, error,
-) {
+) (*clusterController, error) {
 	slog.Info("Starting WireGuard network.")
 	wgnet, err := network.NewWireGuardNetwork()
 	if err != nil {
@@ -69,7 +68,7 @@ func newNetworkController(
 	}
 	endpointChanges := wgnet.WatchEndpoints()
 
-	return &networkController{
+	return &clusterController{
 		state:           state,
 		store:           store,
 		wgnet:           wgnet,
@@ -77,59 +76,55 @@ func newNetworkController(
 		server:          server,
 		corroService:    corroService,
 		dockerCli:       dockerCli,
-		caddyfileCtrl:   caddyfileCtrl,
+		dockerReady:     dockerReady,
+		caddyconfigCtrl: caddyfileCtrl,
 		dnsServer:       dnsServer,
 		dnsResolver:     dnsResolver,
-		networkReady:    networkReady,
 	}, nil
 }
 
-func (nc *networkController) Run(ctx context.Context) error {
+func (cc *clusterController) Run(ctx context.Context) error {
 	if err := firewall.ConfigureIptablesChains(); err != nil {
 		return fmt.Errorf("configure iptables chains: %w", err)
 	}
 
-	if err := nc.wgnet.Configure(*nc.state.Network); err != nil {
+	if err := cc.wgnet.Configure(*cc.state.Network); err != nil {
 		return fmt.Errorf("configure WireGuard network: %w", err)
 	}
 	slog.Info("WireGuard network configured.")
 
-	if nc.corroService.Running() {
+	if cc.corroService.Running() {
 		// Corrosion service was running before the WireGuard network was configured so we need to restart it.
 		slog.Info("Restarting corrosion service to apply new configuration with WireGuard network.")
-		if err := nc.corroService.Restart(ctx); err != nil {
+		if err := cc.corroService.Restart(ctx); err != nil {
 			return fmt.Errorf("restart corrosion service: %w", err)
 		}
 	} else {
 		slog.Info("Starting corrosion service.")
-		if err := nc.corroService.Start(ctx); err != nil {
+		if err := cc.corroService.Start(ctx); err != nil {
 			return fmt.Errorf("start corrosion service: %w", err)
 		}
 	}
-	// TODO: Figure out if we need to manually stop the corrosion service when the context is done or just
-	//  rely on systemd to handle service dependencies on its own.
 
 	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))
+	apiAddr := net.JoinHostPort(cc.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
-		},
-	)
+	errGroup.Go(func() error {
+		slog.Info("Starting network API server.", "addr", apiAddr)
+		if err := cc.server.Serve(listener); err != nil {
+			return fmt.Errorf("network API server failed: %w", err)
+		}
+		return nil
+	})
 
 	errGroup.Go(func() error {
 		slog.Info("Starting embedded DNS resolver.")
-		if err := nc.dnsResolver.Run(ctx); err != nil {
+		if err := cc.dnsResolver.Run(ctx); err != nil {
 			return fmt.Errorf("embedded DNS resolver failed: %w", err)
 		}
 		return nil
@@ -137,7 +132,7 @@ func (nc *networkController) Run(ctx context.Context) error {
 
 	errGroup.Go(func() error {
 		slog.Info("Starting embedded DNS server.")
-		if err := nc.dnsServer.Run(ctx); err != nil {
+		if err := cc.dnsServer.Run(ctx); err != nil {
 			return fmt.Errorf("embedded DNS server failed: %w", err)
 		}
 		return nil
@@ -145,13 +140,13 @@ func (nc *networkController) Run(ctx context.Context) error {
 
 	// Setup Docker network and synchronise containers to the cluster store.
 	errGroup.Go(func() error {
-		return nc.prepareAndWatchDocker(ctx)
+		return cc.prepareAndWatchDocker(ctx)
 	})
 
 	// 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.handleMachineChanges(ctx); err != nil {
+		if err := cc.handleMachineChanges(ctx); err != nil {
 			return fmt.Errorf("handle new machines: %w", err)
 		}
 		return nil
@@ -161,24 +156,24 @@ func (nc *networkController) Run(ctx context.Context) error {
 	errGroup.Go(func() error {
 		for {
 			select {
-			case e, ok := <-nc.endpointChanges:
+			case e, ok := <-cc.endpointChanges:
 				if !ok {
 					// The channel was closed, stop watching for changes.
-					nc.endpointChanges = nil
+					cc.endpointChanges = nil
 					return nil
 				}
 
-				nc.state.mu.Lock()
-				for i := range nc.state.Network.Peers {
-					if nc.state.Network.Peers[i].PublicKey.Equal(e.PublicKey) {
-						nc.state.Network.Peers[i].Endpoint = &e.Endpoint
+				cc.state.mu.Lock()
+				for i := range cc.state.Network.Peers {
+					if cc.state.Network.Peers[i].PublicKey.Equal(e.PublicKey) {
+						cc.state.Network.Peers[i].Endpoint = &e.Endpoint
 						break
 					}
 				}
-				if err := nc.state.Save(); err != nil {
+				if err := cc.state.Save(); err != nil {
 					slog.Error("Failed to save machine state.", "err", err)
 				}
-				nc.state.mu.Unlock()
+				cc.state.mu.Unlock()
 
 				slog.Debug("Preserved endpoint change in the machine state.",
 					"public_key", e.PublicKey, "endpoint", e.Endpoint)
@@ -189,49 +184,56 @@ func (nc *networkController) Run(ctx context.Context) error {
 	})
 
 	errGroup.Go(func() error {
-		if err := nc.wgnet.Run(ctx); err != nil {
+		if err := cc.wgnet.Run(ctx); err != nil {
 			return fmt.Errorf("WireGuard network failed: %w", err)
 		}
 		return nil
 	})
 
 	errGroup.Go(func() error {
-		slog.Info("Starting Caddyconfig controller.")
-		if err := nc.caddyfileCtrl.Run(ctx); err != nil {
-			//goland:noinspection GoErrorStringFormat
-			return fmt.Errorf("Caddyconfig controller failed: %w", err)
+		slog.Info("Starting caddyconfig controller.")
+		if err := cc.caddyconfigCtrl.Run(ctx); err != nil {
+			return fmt.Errorf("caddyconfig controller 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
-	})
+	<-ctx.Done()
+	slog.Info("Stopping network API server.")
+	// TODO: implement timeout for graceful shutdown.
+	cc.server.GracefulStop()
+	slog.Info("Network API server stopped.")
 
-	return errGroup.Wait()
+	// Wait for all controllers to finish.
+	err = errGroup.Wait()
+
+	// It's safe to stop the Corrosion service after the controllers depending on it and API server are stopped.
+	if corroErr := cc.corroService.Stop(ctx); corroErr != nil {
+		slog.Error("Failed to stop corrosion service.", "err", corroErr)
+		err = errors.Join(err, fmt.Errorf("stop corrosion service: %w", corroErr))
+	} else {
+		slog.Info("Corrosion service stopped.")
+	}
+
+	return err
 }
 
 // prepareAndWatchDocker configures the Docker network and watches local Docker containers to sync them
 // to the cluster store.
-func (nc *networkController) prepareAndWatchDocker(ctx context.Context) error {
-	manager := docker.NewManager(nc.dockerCli, nc.state.ID, nc.store)
+func (cc *clusterController) prepareAndWatchDocker(ctx context.Context) error {
+	manager := docker.NewManager(cc.dockerCli, cc.state.ID, cc.store)
 	if err := manager.WaitDaemonReady(ctx); err != nil {
 		return fmt.Errorf("wait for Docker daemon: %w", err)
 	}
 
-	if err := manager.EnsureUncloudNetwork(ctx, nc.state.Network.Subnet, nc.dnsServer.ListenAddr()); err != nil {
+	if err := manager.EnsureUncloudNetwork(ctx, cc.state.Network.Subnet, cc.dnsServer.ListenAddr()); err != nil {
 		return fmt.Errorf("ensure Docker network: %w", err)
 	}
 	slog.Info("Docker network configured.")
 
-	// Signal that the Docker network is ready for containers
-	close(nc.networkReady)
+	// Signal that Docker is ready for containers.
+	close(cc.dockerReady)
 
 	slog.Info("Watching Docker containers and syncing them to cluster store.")
 	// Retry to watch and sync containers until the context is done.
@@ -259,7 +261,7 @@ func (nc *networkController) prepareAndWatchDocker(ctx context.Context) error {
 
 // handleMachineChanges subscribes to machine changes in the cluster and reconfigures the network peers accordingly
 // when changes occur.
-func (nc *networkController) handleMachineChanges(ctx context.Context) error {
+func (cc *clusterController) handleMachineChanges(ctx context.Context) error {
 	for {
 		// Retry to subscribe to machine changes indefinitely until the context is done.
 		boff := backoff.WithContext(backoff.NewExponentialBackOff(
@@ -274,7 +276,7 @@ func (nc *networkController) handleMachineChanges(ctx context.Context) error {
 			err      error
 		)
 		subscribe := func() error {
-			if machines, changes, err = nc.store.SubscribeMachines(ctx); err != nil {
+			if machines, changes, err = cc.store.SubscribeMachines(ctx); err != nil {
 				slog.Info("Failed to subscribe to machine changes, retrying.", "err", err)
 			}
 			return err
@@ -292,7 +294,7 @@ func (nc *networkController) handleMachineChanges(ctx context.Context) error {
 		// completes. Skip configuration now and apply it when the store changes are received.
 		if len(machines) > 0 {
 			slog.Info("Reconfiguring network peers with the current machines.", "machines", len(machines))
-			if err = nc.configurePeers(machines); err != nil {
+			if err = cc.configurePeers(machines); err != nil {
 				slog.Error("Failed to configure peers.", "err", err)
 			}
 		}
@@ -304,11 +306,11 @@ func (nc *networkController) handleMachineChanges(ctx context.Context) error {
 			//  be reworked as well.
 			case <-changes:
 				slog.Info("Cluster machines changed, reconfiguring network peers.")
-				if machines, err = nc.store.ListMachines(ctx); err != nil {
+				if machines, err = cc.store.ListMachines(ctx); err != nil {
 					slog.Error("Failed to list machines.", "err", err)
 					continue
 				}
-				if err = nc.configurePeers(machines); err != nil {
+				if err = cc.configurePeers(machines); err != nil {
 					slog.Error("Failed to configure peers.", "err", err)
 				}
 			case <-ctx.Done():
@@ -318,23 +320,23 @@ func (nc *networkController) handleMachineChanges(ctx context.Context) error {
 	}
 }
 
-func (nc *networkController) configurePeers(machines []*pb.MachineInfo) error {
+func (cc *clusterController) configurePeers(machines []*pb.MachineInfo) error {
 	if len(machines) == 0 {
 		return fmt.Errorf("no machines to configure peers")
 	}
 
-	nc.state.mu.RLock()
-	currentPeerEndpoints := make(map[string]*netip.AddrPort, len(nc.state.Network.Peers))
-	for _, p := range nc.state.Network.Peers {
+	cc.state.mu.RLock()
+	currentPeerEndpoints := make(map[string]*netip.AddrPort, len(cc.state.Network.Peers))
+	for _, p := range cc.state.Network.Peers {
 		currentPeerEndpoints[p.PublicKey.String()] = p.Endpoint
 	}
-	nc.state.mu.RUnlock()
+	cc.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 {
+		if m.Id == cc.state.ID {
 			continue
 		}
 		if err := m.Network.Validate(); err != nil {
@@ -367,17 +369,17 @@ func (nc *networkController) configurePeers(machines []*pb.MachineInfo) error {
 	}
 
 	// 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()
+	cc.state.mu.Lock()
+	cc.state.Network.Peers = peers
+	err := cc.state.Save()
+	cc.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 {
+	cc.state.mu.RLock()
+	defer cc.state.mu.RUnlock()
+	if err = cc.wgnet.Configure(*cc.state.Network); err != nil {
 		return fmt.Errorf("configure network peers: %w", err)
 	}
 	return nil

internal/machine/corroservice/docker.go

@@ -52,6 +52,20 @@ func (s *DockerService) Start(ctx context.Context) error {
 	return s.startNewContainer(ctx)
 }
 
+func (s *DockerService) Stop(ctx context.Context) error {
+	if err := s.Client.ContainerStop(ctx, s.Name, container.StopOptions{}); err != nil {
+		return fmt.Errorf("stop container %q: %w", s.Name, err)
+	}
+	slog.Debug("Corrosion Docker container stopped.", "name", s.Name)
+
+	if err := s.Client.ContainerRemove(ctx, s.Name, container.RemoveOptions{}); err != nil {
+		return fmt.Errorf("remove container %q: %w", s.Name, err)
+	}
+	slog.Debug("Corrosion Docker container removed.", "name", s.Name)
+
+	return nil
+}
+
 func (s *DockerService) Restart(ctx context.Context) error {
 	if err := s.Client.ContainerRestart(ctx, s.Name, container.StopOptions{}); err != nil {
 		return fmt.Errorf("restart container %q: %w", s.Name, err)

internal/machine/corroservice/service.go

@@ -4,6 +4,7 @@ import "context"
 
 type Service interface {
 	Start(ctx context.Context) error
+	Stop(ctx context.Context) error
 	Restart(ctx context.Context) error
 	Running() bool
 }

internal/machine/corroservice/subprocess.go

@@ -1,157 +0,0 @@
-package corroservice
-
-import (
-	"bufio"
-	"context"
-	"fmt"
-	"log/slog"
-	"os/exec"
-	"path/filepath"
-	"sync"
-	"syscall"
-	"time"
-)
-
-const (
-	DefaultCommand = "corrosion"
-	DefaultDataDir = "/var/lib/uncloud/corrosion"
-)
-
-// SubprocessService implements the Service interface by running the service as a subprocess.
-type SubprocessService struct {
-	Command string
-	DataDir string
-
-	cmd         *exec.Cmd
-	running     bool
-	mu          sync.Mutex
-	cancelWatch context.CancelFunc
-}
-
-func DefaultSubprocessService() *SubprocessService {
-	return &SubprocessService{
-		Command: DefaultCommand,
-		DataDir: DefaultDataDir,
-	}
-}
-
-// TODO: maybe stop the process if this ctx is cancelled.
-func (s *SubprocessService) Start(ctx context.Context) error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	if s.running {
-		return nil
-	}
-	return s.startProcess(ctx)
-}
-
-func (s *SubprocessService) Restart(ctx context.Context) error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	if s.running {
-		if err := s.stopProcess(); err != nil {
-			return fmt.Errorf("stop process: %w", err)
-		}
-	}
-	return s.startProcess(ctx)
-}
-
-func (s *SubprocessService) Running() bool {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	return s.running
-}
-
-func (s *SubprocessService) startProcess(ctx context.Context) error {
-	s.cmd = exec.Command(s.Command, "agent", "-c", filepath.Join(s.DataDir, "config.toml"))
-
-	// Redirect stdout and stderr to the logger.
-	stdout, err := s.cmd.StdoutPipe()
-	if err != nil {
-		return fmt.Errorf("create stdout pipe: %w", err)
-	}
-	stderr, err := s.cmd.StderrPipe()
-	if err != nil {
-		return fmt.Errorf("create stderr pipe: %w", err)
-	}
-
-	go func() {
-		scanner := bufio.NewScanner(stdout)
-		for scanner.Scan() {
-			slog.Info("[corrosion]: " + scanner.Text())
-		}
-		// TODO: remove
-		slog.Info("######## corrosion redirect go routine end ########")
-	}()
-
-	go func() {
-		scanner := bufio.NewScanner(stderr)
-		for scanner.Scan() {
-			slog.Error("[corrosion]: " + scanner.Text())
-		}
-	}()
-
-	if err = s.cmd.Start(); err != nil {
-		return fmt.Errorf("start process: %w", err)
-	}
-	s.running = true
-
-	// Watch for process exit to update running status.
-	go func() {
-		if err := s.cmd.Wait(); err != nil {
-			slog.Error("corrosion process exited with error.", "code", s.cmd.ProcessState.ExitCode(), "err", err)
-		}
-
-		s.mu.Lock()
-		s.running = false
-		s.mu.Unlock()
-	}()
-
-	// TODO: figure out the waiting process
-	// Wait for initialization
-	// timer := time.NewTimer(2 * time.Second)
-	// defer timer.Stop()
-
-	//select {
-	////case <-timer.C:
-	////	s.running = true
-	////	return nil
-	//case <-watchCtx.Done():
-	//	return fmt.Errorf("process failed to start")
-	//case <-ctx.Done():
-	//	s.stopProcess()
-	//	return ctx.Err()
-	//}
-	return nil
-}
-
-func (s *SubprocessService) stopProcess() error {
-	if s.cmd == nil || s.cmd.Process == nil {
-		return nil
-	}
-
-	if err := s.cmd.Process.Signal(syscall.SIGTERM); err != nil {
-		return fmt.Errorf("send SIGTERM: %w", err)
-	}
-
-	// Wait up to 5 seconds for graceful shutdown before killing the process.
-	done := make(chan error, 1)
-	go func() {
-		done <- s.cmd.Wait()
-	}()
-
-	select {
-	case <-time.After(5 * time.Second):
-		if err := s.cmd.Process.Kill(); err != nil {
-			return fmt.Errorf("kill process: %w", err)
-		}
-	case err := <-done:
-		if err != nil {
-			return fmt.Errorf("process exited with error: %w", err)
-		}
-	}
-
-	return nil
-}

internal/machine/corroservice/systemd.go

@@ -27,6 +27,15 @@ func (s *SystemdService) Start(ctx context.Context) error {
 	return s.startOrRestart(ctx, "start")
 }
 
+func (s *SystemdService) Stop(ctx context.Context) error {
+	if _, err := exec.Command("systemctl", "stop", s.Unit).Output(); err != nil {
+		return fmt.Errorf("systemctl stop %s: %w", s.Unit, err)
+	}
+	slog.Info("Corrosion systemd service stopped.", "unit", s.Unit)
+
+	return nil
+}
+
 func (s *SystemdService) Restart(ctx context.Context) error {
 	return s.startOrRestart(ctx, "restart")
 }

internal/machine/machine.go

@@ -152,9 +152,10 @@ type Machine struct {
 	initialised chan struct{}
 	// networkReady is signalled when the Docker network is configured and ready for containers.
 	networkReady chan struct{}
-	// networkReadyMu protects networkReady channel operations
-	networkReadyMu sync.RWMutex
+	// stop cancels the Run method context to stop the machine gracefully.
+	stop func()
 
+	clusterCtrl *clusterController
 	// store is the cluster store backed by a distributed Corrosion database.
 	store   *store.Store
 	cluster *cluster.Cluster
@@ -168,6 +169,9 @@ type Machine struct {
 	// It proxies requests to the local or remote machine API servers depending on the request targets
 	// and aggregates responses.
 	localProxyServer *grpc.Server
+
+	// mu protects the Machine from concurrent reads and writes.
+	mu sync.RWMutex
 }
 
 func NewMachine(config *Config) (*Machine, error) {
@@ -258,10 +262,6 @@ func NewMachine(config *Config) (*Machine, error) {
 
 	if m.Initialised() {
 		m.initialised <- struct{}{}
-	} else {
-		// For non-initialized machines, signal network is ready immediately
-		// since there's no cluster network to set up
-		close(m.networkReady)
 	}
 
 	return m, nil
@@ -299,6 +299,9 @@ func (m *Machine) IP() netip.Addr {
 }
 
 func (m *Machine) Run(ctx context.Context) error {
+	// Create a cancellable context for the Run method to allow stopping the machine gracefully.
+	ctx, m.stop = context.WithCancel(ctx)
+
 	// Docker dependency is essential for the machine to function. Block until it's ready.
 	if err := docker.WaitDaemonReady(ctx, m.config.DockerClient); err != nil {
 		return fmt.Errorf("wait for Docker daemon: %w", err)
@@ -306,7 +309,7 @@ func (m *Machine) Run(ctx context.Context) error {
 
 	// Configure and start the corrosion service on the loopback if the machine is not initialised as a cluster
 	// member. This provides the store required for the machine to initialise a new cluster on it. Once the machine
-	// is initialised, the corrosion service is managed by the networkController.
+	// is initialised, the corrosion service is managed by the clusterController.
 	if !m.Initialised() {
 		if err := m.configureCorrosion(); err != nil {
 			return fmt.Errorf("configure corrosion service: %w", err)
@@ -326,153 +329,114 @@ func (m *Machine) Run(ctx context.Context) error {
 	if err != nil {
 		return fmt.Errorf("listen machine API unix socket %q: %w", m.config.MachineSockPath, err)
 	}
-	errGroup.Go(
-		func() error {
-			slog.Info("Starting local machine API server.", "path", m.config.MachineSockPath)
-			if err := m.localMachineServer.Serve(machineListener); err != nil {
-				return fmt.Errorf("local machine API server failed: %w", err)
-			}
-			return nil
-		},
-	)
+	errGroup.Go(func() error {
+		slog.Info("Starting local machine API server.", "path", m.config.MachineSockPath)
+		if err := m.localMachineServer.Serve(machineListener); err != nil {
+			return fmt.Errorf("local machine API server failed: %w", err)
+		}
+		return nil
+	})
 
 	// Start the local API proxy server.
 	proxyListener, err := listenUnixSocket(m.config.UncloudSockPath)
 	if err != nil {
 		return fmt.Errorf("listen API proxy unix socket %q: %w", m.config.UncloudSockPath, err)
 	}
-	errGroup.Go(
-		func() error {
-			slog.Info("Starting local API proxy server.", "path", m.config.UncloudSockPath)
-			if err := m.localProxyServer.Serve(proxyListener); err != nil {
-				return fmt.Errorf("local API proxy server failed: %w", err)
-			}
-			return nil
-		},
-	)
+	errGroup.Go(func() error {
+		slog.Info("Starting local API proxy server.", "path", m.config.UncloudSockPath)
+		if err := m.localProxyServer.Serve(proxyListener); err != nil {
+			return fmt.Errorf("local API proxy server failed: %w", err)
+		}
+		return nil
+	})
 	// Signal that the machine is ready.
 	close(m.started)
 
-	// Control loop for managing components that depend on the machine being initialised as a cluster member.
-	errGroup.Go(
-		func() error {
-			if !m.Initialised() {
-				slog.Info(
-					"Waiting for the machine to be initialised as a member of a cluster " +
-						"to start the network controller.",
-				)
-			}
+	// Wait for the machine to be initialised as a member of a cluster and run the cluster controller.
+	errGroup.Go(func() error {
+		if !m.Initialised() {
+			slog.Info(
+				"Waiting for the machine to be initialised as a member of a cluster to start the cluster controller.",
+			)
+		}
+		<-m.initialised
 
-			var ctrl *networkController
-			// Error channel for communicating the termination of the network controller.
-			errCh := make(chan error)
+		m.cluster.UpdateMachineID(m.state.ID)
 
-			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:
-					var err error
+		// Ensure the corrosion config is up to date, including a new gossip address if the machine
+		// has just joined a cluster.
+		if err := m.configureCorrosion(); err != nil {
+			return fmt.Errorf("configure corrosion service: %w", err)
+		}
+		slog.Info("Configured corrosion service.", "dir", m.config.CorrosionDir)
 
-					// Reset networkReady channel for the new cluster configuration
-					m.networkReadyMu.Lock()
-					m.networkReady = make(chan struct{})
-					m.networkReadyMu.Unlock()
+		slog.Info("Starting cluster controller.")
+		// Update the proxy director's local address to the machine's management IP address, allowing
+		// the proxy to identify which requests should be proxied to the local machine API server.
+		m.proxyDirector.UpdateLocalAddress(m.state.Network.ManagementIP.String())
+		proxyServer := grpc.NewServer(
+			grpc.ForceServerCodecV2(proxy.Codec()),
+			grpc.UnknownServiceHandler(
+				proxy.TransparentHandler(m.proxyDirector.Director),
+			),
+		)
 
-					m.cluster.UpdateMachineID(m.state.ID)
+		// Create a new caddyconfig controller for managing the Caddy reverse proxy configuration.
+		// It will also serve the current machine ID at /.uncloud-verify to verify Caddy reachability.
+		caddyconfigCtrl, err := caddyconfig.NewController(m.store, m.config.CaddyConfigPath, m.state.ID)
+		if err != nil {
+			return fmt.Errorf("create caddyconfig controller: %w", err)
+		}
 
-					// Ensure the corrosion config is up to date, including a new gossip address if the machine
-					// has just joined a cluster.
-					if err = m.configureCorrosion(); err != nil {
-						return fmt.Errorf("configure corrosion service: %w", err)
-					}
-					slog.Info("Configured corrosion service.", "dir", m.config.CorrosionDir)
+		dnsResolver := dns.NewClusterResolver(m.store)
+		dnsServer, err := dns.NewServer(m.IP(), dnsResolver, m.config.DNSUpstreams)
+		if err != nil {
+			return fmt.Errorf("create embedded DNS server: %w", err)
+		}
 
-					slog.Info("Starting network controller.")
-					// Update the proxy director's local address to the machine's management IP address, allowing
-					// the proxy to identify which requests should be proxied to the local machine API server.
-					m.proxyDirector.UpdateLocalAddress(m.state.Network.ManagementIP.String())
-					proxyServer := grpc.NewServer(
-						grpc.ForceServerCodecV2(proxy.Codec()),
-						grpc.UnknownServiceHandler(
-							proxy.TransparentHandler(m.proxyDirector.Director),
-						),
-					)
+		m.mu.Lock()
+		m.clusterCtrl, err = newClusterController(
+			m.state,
+			m.store,
+			proxyServer,
+			m.config.CorrosionService,
+			m.config.DockerClient,
+			m.networkReady,
+			caddyconfigCtrl,
+			dnsServer,
+			dnsResolver,
+		)
+		m.mu.Unlock()
+		if err != nil {
+			return fmt.Errorf("initialise cluster controller: %w", err)
+		}
 
-					// Create a new Caddyfile controller for managing the Caddy reverse proxy configuration.
-					// It will also serve the current machine ID at /.uncloud-verify to verify Caddy reachability.
-					caddyfileCtrl, err := caddyconfig.NewController(m.store, m.config.CaddyConfigPath, m.state.ID)
-					if err != nil {
-						return fmt.Errorf("create Caddyfile controller: %w", err)
-					}
+		if err = m.clusterCtrl.Run(ctx); err != nil {
+			return fmt.Errorf("run cluster controller: %w", err)
+		}
+		slog.Info("Cluster controller stopped.")
 
-					dnsResolver := dns.NewClusterResolver(m.store)
-					dnsServer, err := dns.NewServer(m.IP(), dnsResolver, m.config.DNSUpstreams)
-					if err != nil {
-						return fmt.Errorf("create embedded DNS server: %w", err)
-					}
-
-					ctrl, err = newNetworkController(
-						m.state,
-						m.store,
-						proxyServer,
-						m.config.CorrosionService,
-						m.config.DockerClient,
-						caddyfileCtrl,
-						dnsServer,
-						dnsResolver,
-						m.networkReady,
-					)
-					if err != nil {
-						return fmt.Errorf("initialise network controller: %w", err)
-					}
-
-					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
-					}
-					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
-				}
-			}
-		},
-	)
+		return nil
+	})
 
 	// Shutdown goroutine.
-	errGroup.Go(
-		func() error {
-			<-ctx.Done()
-			slog.Info("Stopping local machine API server.")
-			// TODO: implement timeout for graceful shutdown.
-			m.localMachineServer.GracefulStop()
-			slog.Info("Local machine API server stopped.")
+	errGroup.Go(func() error {
+		<-ctx.Done()
+		slog.Info("Stopping local machine API server.")
+		// TODO: implement timeout for graceful shutdown.
+		m.localMachineServer.GracefulStop()
+		slog.Info("Local machine API server stopped.")
 
-			slog.Info("Stopping local API proxy server.")
-			// TODO: implement timeout for graceful shutdown.
-			m.localProxyServer.GracefulStop()
-			// Close the proxy director to close all backend connections.
-			m.proxyDirector.Close()
-			slog.Info("Local API proxy server stopped.")
+		slog.Info("Stopping local API proxy server.")
+		// TODO: implement timeout for graceful shutdown.
+		m.localProxyServer.GracefulStop()
+		// Close the proxy director to close all backend connections.
+		m.proxyDirector.Close()
+		slog.Info("Local API proxy server stopped.")
 
-			m.config.DockerClient.Close()
-			return nil
-		},
-	)
+		m.config.DockerClient.Close()
+		return nil
+	})
 
 	return errGroup.Wait()
 }
@@ -798,13 +762,10 @@ func (m *Machine) Inspect(_ context.Context, _ *emptypb.Empty) (*pb.MachineInfo,
 func (m *Machine) IsNetworkReady() bool {
 	if !m.Initialised() {
 		// If machine is not initialized, there's no network to check
-		return true
+		return false
 	}
 
 	// Check if network is ready by checking if the networkReady channel has been closed
-	m.networkReadyMu.RLock()
-	defer m.networkReadyMu.RUnlock()
-
 	select {
 	case <-m.networkReady:
 		return true
@@ -821,34 +782,43 @@ func (m *Machine) WaitForNetworkReady(ctx context.Context) error {
 		return nil
 	}
 
-	// Get a copy of the channel to wait on
-	m.networkReadyMu.RLock()
-	networkReady := m.networkReady
-	m.networkReadyMu.RUnlock()
-
 	// Wait for network to be ready or context to be cancelled
 	select {
-	case <-networkReady:
+	case <-m.networkReady:
 		return nil
 	case <-ctx.Done():
 		return ctx.Err()
 	}
 }
 
-// Reset restores the machine to a clean state, removing all cluster-related сonfiguration and data and scheduling
+// Reset restores the machine to a clean state, removing all cluster-related configuration and data and scheduling
 // a graceful shutdown. The uncloud daemon will restart the machine if managed by systemd.
 func (m *Machine) Reset(ctx context.Context, _ *pb.ResetRequest) (*emptypb.Empty, error) {
+	if !m.Initialised() {
+		return nil, nil
+	}
+
 	slog.Info("Resetting machine to a clean state.")
 
 	// TODO: stop and remove all managed service containers.
 	// TODO: check if the request is coming from the unix or network socket. For the network socket, the reset should
 	//  be called in a separate goroutine to avoid blocking the RPC response.
-	// TODO: stop the network controller
-	// TODO: implement and call Cleanup on the network controller to remove Docker network, WG interface, iptables
-	//  rules, corrosion state, ?stop corrosion service.
-	// TODO: stop the machine and remove the machine.json state. The daemon should restart it to a clean state.
 
-	return &emptypb.Empty{}, status.Error(codes.Unimplemented, "reset machine is not implemented yet")
+	// TODO: Stop the machine asynchronously. The gRPC servers will wait for this request to complete before stopping.
+	go func() {
+		m.stop()
+		// TODO: wait for the cluster controller to stop.
+		// TODO: Cleanup cluster controller (WG network, Docker network, iptables rules, etc.)
+
+		// TODO: uncomment after testing all other cleanup steps.
+		//if err := os.RemoveAll(m.config.DataDir); err != nil {
+		//	slog.Error("Failed to remove data directory storing persistent machine state.",
+		//		"path", m.config.DataDir, "err", err)
+		//	return nil, status.Errorf(codes.Internal, "remove data directory on machine '%s': %v", m.config.DataDir, err)
+		//}
+		//slog.Info("Removed data directory storing persistent machine state.", "path", m.config.DataDir)
+	}()
+	return &emptypb.Empty{}, nil
 }
 
 // InspectService returns detailed information about a service and its containers stored in the cluster store.