uncloud/internal/machine/cluster/cluster.go

package cluster

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/emptypb"
	"log/slog"
	"net/netip"
	"time"
	"uncloud/internal/machine/api/pb"
	"uncloud/internal/machine/network"
	"uncloud/internal/machine/store"
	"uncloud/internal/secret"
)

type Cluster struct {
	pb.UnimplementedClusterServer

	state *State
	store *store.Store

	// TODO: temporary channel until the state is replaced with networkDB.
	newMachinesCh chan *pb.MachineInfo
}

func NewCluster(state *State, store *store.Store) *Cluster {
	return &Cluster{
		state:         state,
		store:         store,
		newMachinesCh: make(chan *pb.MachineInfo, 1),
	}
}

func (c *Cluster) Init(ctx context.Context, network netip.Prefix) error {
	initialised, err := c.Initialised(ctx)
	if err != nil {
		return err
	}
	if initialised {
		return fmt.Errorf("cluster already initialized")
	}

	if err = c.store.Put(ctx, "network", network.String()); err != nil {
		return fmt.Errorf("put network to store: %w", err)
	}
	if err = c.store.Put(ctx, "created_at", time.Now().UTC().Format(time.RFC3339)); err != nil {
		return fmt.Errorf("put created_at to store: %w", err)
	}
	return nil
}

func (c *Cluster) Initialised(ctx context.Context) (bool, error) {
	var createdAt string
	if err := c.store.Get(ctx, "created_at", &createdAt); err != nil {
		if errors.Is(err, store.ErrKeyNotFound) {
			return false, nil
		}
		return false, fmt.Errorf("get created_at from store: %w", err)
	}
	return true, nil
}

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 *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")
	}
	if network == nil {
		return fmt.Errorf("network not set")
	}
	c.state.State.Network = network
	return 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 c.state == nil {
		return nil, status.Error(codes.FailedPrecondition, "cluster is not initialized")
	}

	if err := req.Validate(); err != nil {
		return nil, err
	}
	if len(req.Network.Endpoints) == 0 {
		return nil, status.Error(codes.InvalidArgument, "endpoints not set")
	}

	machines := c.state.State.Machines
	allocatedSubnets := make([]netip.Prefix, len(machines))
	var err error
	i := 0
	for _, m := range machines {
		if req.Name != "" && m.Name == req.Name {
			return nil, status.Errorf(codes.AlreadyExists, "machine with name %q already exists", req.Name)
		}
		if m.Network.ManagementIp != nil && req.Network.ManagementIp != nil &&
			m.Network.ManagementIp.Equal(req.Network.ManagementIp) {
			manageIP, _ := req.Network.ManagementIp.ToAddr()
			return nil, status.Errorf(codes.AlreadyExists, "machine with management IP %q already exists", manageIP)
		}
		if bytes.Equal(m.Network.PublicKey, req.Network.PublicKey) {
			publicKey := secret.Secret(m.Network.PublicKey)
			return nil, status.Errorf(codes.AlreadyExists, "machine with public key %q already exists", publicKey)
		}

		if allocatedSubnets[i], err = m.Network.Subnet.ToPrefix(); err != nil {
			return nil, err
		}
		i++
	}

	mid, err := NewMachineID()
	if err != nil {
		return nil, status.Errorf(codes.Internal, "generate machine ID: %v", err)
	}
	manageIP := req.Network.ManagementIp
	if manageIP == nil {
		manageIP = pb.NewIP(network.ManagementIP(req.Network.PublicKey))
	}
	m := &pb.MachineInfo{
		Id:   mid,
		Name: req.Name,
		Network: &pb.NetworkConfig{
			ManagementIp: manageIP,
			PublicKey:    req.Network.PublicKey,
		},
	}
	if m.Name == "" {
		m.Name, err = NewRandomMachineName()
		if err != nil {
			return nil, status.Errorf(codes.Internal, "generate machine name: %v", err)
		}
	}

	clusterNetwork, err := c.Network()
	if err != nil {
		return nil, status.Errorf(codes.Internal, "get cluster network: %v", err)
	}
	ipam, err := NewIPAMWithAllocated(clusterNetwork, allocatedSubnets)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "create IPAM manager: %v", err)
	}
	subnet, err := ipam.AllocateSubnetLen(network.DefaultSubnetBits)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "allocate subnet for machine: %v", err)
	}
	m.Network.Subnet = pb.NewIPPrefix(subnet)

	// TODO: announce the new machine to the cluster network using Serf and achieve consensus.
	mCopy := proto.Clone(m).(*pb.MachineInfo)
	c.state.State.Machines[mCopy.Id] = mCopy
	// Store machine endpoints in a separate collection to allow modifying them with limited permissions.
	c.state.State.Endpoints[mCopy.Id] = &pb.MachineEndpoints{
		Id:        mCopy.Id,
		Endpoints: req.Network.Endpoints,
	}
	if err = c.state.Save(); err != nil {
		return nil, status.Errorf(codes.Internal, "save state: %v", err)
	}
	slog.Info("Machine added to the cluster.", "id", m.Id, "name", m.Name)

	// Include the machine endpoints in the response.
	m.Network.Endpoints = req.Network.Endpoints

	// TODO: notify all cluster machines 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
}

func (c *Cluster) ListMachines(ctx context.Context, _ *emptypb.Empty) (*pb.ListMachinesResponse, error) {
	if c.state == nil {
		return nil, status.Error(codes.FailedPrecondition, "cluster is not initialized")
	}
	// TODO: consider creating MachineInfo type that pb.MachineInfo is mapped to always carry valid data to reduce
	//  error handling caused by conversions and nil pointers.
	machines := make([]*pb.MachineInfo, 0, len(c.state.State.Machines))
	for _, sm := range c.state.State.Machines {
		m := proto.Clone(sm).(*pb.MachineInfo)
		m.Network.Endpoints = c.state.State.Endpoints[m.Id].Endpoints
		machines = append(machines, m)
	}
	return &pb.ListMachinesResponse{Machines: machines}, nil
}

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)
	}
	return &pb.ListMachineEndpointsResponse{Endpoints: endpoints}, nil
}

func (c *Cluster) AddUser(user *pb.User) error {
	c.state.State.Users = append(c.state.State.Users, user)
	return c.state.Save()
}

func (c *Cluster) ListUsers() []*pb.User {
	return c.state.State.Users
}