chore: refactor schedulers to work with cluster state snapshot, VolumeScheduler updates state with scheduled volumes

pkg/client/compose/deploy.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/compose-spec/compose-go/v2/graph"
 	"github.com/compose-spec/compose-go/v2/types"
+	"github.com/docker/docker/api/types/volume"
 	"github.com/psviderski/uncloud/pkg/api"
 	"github.com/psviderski/uncloud/pkg/client/deploy"
 	"github.com/psviderski/uncloud/pkg/client/deploy/scheduler"
@@ -23,10 +24,16 @@ type Deployment struct {
 	Client       Client
 	Project      *types.Project
 	SpecResolver *deploy.ServiceSpecResolver
+	state        *scheduler.ClusterState
 	plan         *deploy.SequenceOperation
 }
 
 func NewDeployment(ctx context.Context, cli Client, project *types.Project) (*Deployment, error) {
+	state, err := scheduler.InspectClusterState(ctx, cli)
+	if err != nil {
+		return nil, fmt.Errorf("inspect cluster state: %w", err)
+	}
+
 	domain, err := cli.GetDomain(ctx)
 	if err != nil && !errors.Is(err, api.ErrNotFound) {
 		return nil, fmt.Errorf("get cluster domain: %w", err)
@@ -41,6 +48,7 @@ func NewDeployment(ctx context.Context, cli Client, project *types.Project) (*De
 		Client:       cli,
 		Project:      project,
 		SpecResolver: resolver,
+		state:        state,
 	}, nil
 }
 
@@ -67,7 +75,7 @@ func (d *Deployment) Plan(ctx context.Context) (deploy.SequenceOperation, error)
 	}
 
 	// Check external volumes and plan the creation of missing volumes before deploying services.
-	volumeOps, err := d.planVolumes(ctx, serviceSpecs)
+	volumeOps, err := d.planVolumes(serviceSpecs)
 	if err != nil {
 		return plan, err
 	}
@@ -77,7 +85,8 @@ func (d *Deployment) Plan(ctx context.Context) (deploy.SequenceOperation, error)
 
 	for _, spec := range serviceSpecs {
 		// TODO: properly handle depends_on conditions in the service deployment plan as the first operation.
-		deployment := deploy.NewDeployment(d.Client, spec, nil)
+		// Pass the update cluster state with scheduled volumes to the deployment.
+		deployment := deploy.NewDeployment(d.Client, spec, &deploy.RollingStrategy{State: d.state})
 		servicePlan, err := deployment.Plan(ctx)
 		if err != nil {
 			return plan, fmt.Errorf("create deployment plan for service '%s': %w", spec.Name, err)
@@ -104,21 +113,19 @@ func (d *Deployment) ServiceSpec(name string) (api.ServiceSpec, error) {
 }
 
 // PlanVolumes checks if the external volumes exist and plans the creation of missing volumes.
-func (d *Deployment) planVolumes(
-	ctx context.Context, serviceSpecs []api.ServiceSpec,
-) ([]*deploy.CreateVolumeOperation, error) {
+func (d *Deployment) planVolumes(serviceSpecs []api.ServiceSpec) ([]*deploy.CreateVolumeOperation, error) {
 	if len(d.Project.Volumes) == 0 {
 		// No volumes to check or create.
 		return nil, nil
 	}
 
-	if err := d.checkExternalVolumesExist(ctx); err != nil {
+	if err := d.checkExternalVolumesExist(); err != nil {
 		return nil, err
 	}
 
 	// TODO: The scheduler should ideally work with the resolved service specs to correctly identify eligible machines.
 	//  Figure out where the best place to resolve the specs is.
-	volumeScheduler, err := scheduler.NewVolumeSchedulerWithClient(ctx, d.Client, serviceSpecs)
+	volumeScheduler, err := scheduler.NewVolumeScheduler(d.state, serviceSpecs)
 	if err != nil {
 		return nil, fmt.Errorf("init volume scheduler: %w", err)
 	}
@@ -142,7 +149,7 @@ func (d *Deployment) planVolumes(
 }
 
 // checkExternalVolumesExist checks that all external volumes exist in the cluster.
-func (d *Deployment) checkExternalVolumesExist(ctx context.Context) error {
+func (d *Deployment) checkExternalVolumesExist() error {
 	var externalNames []string
 	for _, v := range d.Project.Volumes {
 		if v.External {
@@ -150,15 +157,12 @@ func (d *Deployment) checkExternalVolumesExist(ctx context.Context) error {
 		}
 	}
 
-	volumes, err := d.Client.ListVolumes(ctx, &api.VolumeFilter{Names: externalNames})
-	if err != nil {
-		return fmt.Errorf("list volumes: %w", err)
-	}
-
 	var notFound []string
 	for _, name := range externalNames {
-		if !slices.ContainsFunc(volumes, func(vol api.MachineVolume) bool {
-			return vol.Volume.Name == name
+		if !slices.ContainsFunc(d.state.Machines, func(m *scheduler.Machine) bool {
+			return slices.ContainsFunc(m.Volumes, func(vol volume.Volume) bool {
+				return vol.Name == name
+			})
 		}) {
 			notFound = append(notFound, fmt.Sprintf("'%s'", name))
 		}

pkg/client/deploy/scheduler/constraint.go

@@ -1,9 +1,11 @@
 package scheduler
 
 import (
+	"reflect"
 	"slices"
 	"strings"
 
+	"github.com/docker/docker/api/types/mount"
 	"github.com/docker/docker/api/types/volume"
 	"github.com/psviderski/uncloud/pkg/api"
 )
@@ -73,7 +75,7 @@ type VolumesConstraint struct {
 }
 
 // Evaluate determines if a machine has all the required volumes.
-// Returns true if all required volumes exist on the machine or if there are no required volumes.
+// Returns true if all required volumes exist or scheduled on the machine or if there are no required volumes.
 func (c *VolumesConstraint) Evaluate(machine *Machine) bool {
 	if len(c.Volumes) == 0 {
 		return true
@@ -84,9 +86,35 @@ func (c *VolumesConstraint) Evaluate(machine *Machine) bool {
 			continue
 		}
 
-		// TODO: should we check the volume driver to be local or any matched volume by name is ok?
-		if !slices.ContainsFunc(machine.Volumes, func(vol volume.Volume) bool {
-			return vol.Name == v.DockerVolumeName()
+		// Check if the required volume already exists on the machine.
+		if slices.ContainsFunc(machine.Volumes, func(vol volume.Volume) bool {
+			if v.DockerVolumeName() == vol.Name {
+				return v.MatchesDockerVolume(vol)
+			}
+			return false
+		}) {
+			continue
+		}
+
+		// Check if the required volume has been scheduled on the machine. The driver names and options must match.
+		if !slices.ContainsFunc(machine.ScheduledVolumes, func(scheduled api.VolumeSpec) bool {
+			if v.DockerVolumeName() != scheduled.DockerVolumeName() {
+				return false
+			}
+
+			// The volume spec with an empty driver can mount the volume that matches the name no matter the driver.
+			if v.VolumeOptions.Driver == nil {
+				return true
+			}
+
+			// If the driver is specified in the spec, the spec's driver and options must match the volume's driver
+			// and options to successfully mount the volume.
+			scheduled = scheduled.SetDefaults()
+			scheduledDriver := scheduled.VolumeOptions.Driver
+			if scheduledDriver == nil {
+				scheduledDriver = &mount.Driver{Name: api.VolumeDriverLocal}
+			}
+			return reflect.DeepEqual(v.VolumeOptions.Driver, scheduledDriver)
 		}) {
 			return false
 		}

pkg/client/deploy/scheduler/service.go

@@ -1,35 +1,24 @@
 package scheduler
 
 import (
-	"context"
 	"errors"
-	"fmt"
 
 	"github.com/psviderski/uncloud/internal/machine/api/pb"
 	"github.com/psviderski/uncloud/pkg/api"
 )
 
 type ServiceScheduler struct {
-	machines    []*Machine
+	state       *ClusterState
 	spec        api.ServiceSpec
 	constraints []Constraint
 }
 
-func NewServiceSchedulerWithClient(ctx context.Context, cli Client, spec api.ServiceSpec) (*ServiceScheduler, error) {
-	machines, err := InspectMachines(ctx, cli)
-	if err != nil {
-		return nil, fmt.Errorf("inspect machines: %w", err)
-	}
-
-	return NewServiceSchedulerWithMachines(machines, spec), nil
-}
-
-// NewServiceSchedulerWithMachines creates a new ServiceScheduler with the given machines and service specification.
-func NewServiceSchedulerWithMachines(machines []*Machine, spec api.ServiceSpec) *ServiceScheduler {
+// NewServiceScheduler creates a new ServiceScheduler with the given cluster state and service specification.
+func NewServiceScheduler(state *ClusterState, spec api.ServiceSpec) *ServiceScheduler {
 	constraints := constraintsFromSpec(spec)
 
 	return &ServiceScheduler{
-		machines:    machines,
+		state:       state,
 		spec:        spec,
 		constraints: constraints,
 	}
@@ -38,7 +27,7 @@ func NewServiceSchedulerWithMachines(machines []*Machine, spec api.ServiceSpec)
 // EligibleMachines returns a list of machines that satisfy all constraints for the next scheduled container.
 func (s *ServiceScheduler) EligibleMachines() ([]*Machine, error) {
 	var available []*Machine
-	for _, machine := range s.machines {
+	for _, machine := range s.state.Machines {
 		if s.evaluateConstraints(machine) {
 			available = append(available, machine)
 		}

pkg/client/deploy/scheduler/state.go

@@ -9,19 +9,26 @@ import (
 	"github.com/psviderski/uncloud/pkg/api"
 )
 
+// ClusterState represents the current and planned state of machines and their resources in the cluster.
+type ClusterState struct {
+	Machines []*Machine
+}
+
+type Machine struct {
+	Info             *pb.MachineInfo
+	Volumes          []volume.Volume
+	ScheduledVolumes []api.VolumeSpec
+}
+
 type Client interface {
 	api.MachineClient
 	api.VolumeClient
 }
 
-type Machine struct {
-	Info    *pb.MachineInfo
-	Volumes []volume.Volume
-}
-
-// InspectMachines retrieves the list of available machines and their details required for scheduling purposes.
-// TODO: refactor to get all the details in one broadcast call to machine API.
-func InspectMachines(ctx context.Context, cli Client) ([]*Machine, error) {
+// InspectClusterState creates a new cluster state by inspecting the machines using the cluster client.
+func InspectClusterState(ctx context.Context, cli Client) (*ClusterState, error) {
+	// TODO: refactor to get all the details in one broadcast call to machine API,
+	//  e.g. InspectMachine with include options.
 	machineMembers, err := cli.ListMachines(ctx, &api.MachineFilter{Available: true})
 	if err != nil {
 		return nil, fmt.Errorf("list machines: %w", err)
@@ -46,5 +53,7 @@ func InspectMachines(ctx context.Context, cli Client) ([]*Machine, error) {
 		machines = append(machines, machine)
 	}
 
-	return machines, nil
+	return &ClusterState{
+		Machines: machines,
+	}, nil
 }

pkg/client/deploy/scheduler/volume.go

@@ -1,7 +1,6 @@
 package scheduler
 
 import (
-	"context"
 	"fmt"
 	"slices"
 	"strings"
@@ -18,8 +17,8 @@ import (
 //   - If a volume already exists on a machine, it must be used instead of creating a new one.
 //   - A missing volume must only be created on one machine.
 type VolumeScheduler struct {
-	// machines is a list of available machines in the cluster.
-	machines []*Machine
+	// state is the current state of machines and their resources in the cluster.
+	state *ClusterState
 	// serviceSpecs is a list of service specifications included in the deployment.
 	serviceSpecs []api.ServiceSpec
 	// volumeSpecs is a map of volume names to their specifications from the service specs in a canonical form.
@@ -32,19 +31,8 @@ type VolumeScheduler struct {
 	existingVolumeMachines map[string]mapset.Set[string]
 }
 
-// NewVolumeSchedulerWithClient creates a new VolumeScheduler with the given cluster client and service specifications.
-func NewVolumeSchedulerWithClient(ctx context.Context, cli Client, specs []api.ServiceSpec) (*VolumeScheduler, error) {
-	machines, err := InspectMachines(ctx, cli)
-	if err != nil {
-		return nil, fmt.Errorf("inspect machines: %w", err)
-	}
-
-	return NewVolumeSchedulerWithMachines(machines, specs)
-}
-
-// NewVolumeSchedulerWithMachines creates a new VolumeScheduler with the given cluster machines
-// and service specifications.
-func NewVolumeSchedulerWithMachines(machines []*Machine, specs []api.ServiceSpec) (*VolumeScheduler, error) {
+// NewVolumeScheduler creates a new VolumeScheduler with the given cluster state and service specifications.
+func NewVolumeScheduler(state *ClusterState, specs []api.ServiceSpec) (*VolumeScheduler, error) {
 	var specsWithVolumes []api.ServiceSpec
 	// Docker volume name -> VolumeSpec.
 	volumeSpecs := make(map[string]api.VolumeSpec)
@@ -90,7 +78,7 @@ func NewVolumeSchedulerWithMachines(machines []*Machine, specs []api.ServiceSpec
 
 	// Validate the configurations of existing volumes on machines don't conflict with the volume specs, for example,
 	// a volume and a spec with the same name don't have different drivers.
-	for _, machine := range machines {
+	for _, machine := range state.Machines {
 		for _, vol := range machine.Volumes {
 			if spec, ok := volumeSpecs[vol.Name]; ok {
 				if !spec.MatchesDockerVolume(vol) {
@@ -107,7 +95,7 @@ func NewVolumeSchedulerWithMachines(machines []*Machine, specs []api.ServiceSpec
 	}
 
 	return &VolumeScheduler{
-		machines:               machines,
+		state:                  state,
 		serviceSpecs:           specsWithVolumes,
 		volumeSpecs:            volumeSpecs,
 		volumeServices:         volumeServices,
@@ -204,6 +192,16 @@ func (s *VolumeScheduler) Schedule() (map[string][]api.VolumeSpec, error) {
 		}
 	}
 
+	// Update the state of the machines with the scheduled volumes.
+	for machineID, volumes := range scheduledVolumes {
+		for _, m := range s.state.Machines {
+			if m.Info.Id == machineID {
+				m.ScheduledVolumes = append(m.ScheduledVolumes, volumes...)
+				break
+			}
+		}
+	}
+
 	return scheduledVolumes, nil
 }
 
@@ -213,7 +211,7 @@ func (s *VolumeScheduler) serviceEligibleMachinesWithoutVolumes(spec api.Service
 	specWithoutVolumes := spec.Clone()
 	specWithoutVolumes.Container.VolumeMounts = nil
 
-	scheduler := NewServiceSchedulerWithMachines(s.machines, specWithoutVolumes)
+	scheduler := NewServiceScheduler(s.state, specWithoutVolumes)
 	machines, err := scheduler.EligibleMachines()
 	if err != nil {
 		return nil, fmt.Errorf("schedule service '%s': %w", spec.Name, err)

pkg/client/deploy/scheduler/volume_test.go

@@ -806,7 +806,10 @@ func TestVolumeScheduler_Schedule(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			scheduler, err := NewVolumeSchedulerWithMachines(tt.machines, tt.serviceSpecs)
+			state := &ClusterState{
+				Machines: tt.machines,
+			}
+			scheduler, err := NewVolumeScheduler(state, tt.serviceSpecs)
 			require.NoError(t, err)
 			result, err := scheduler.Schedule()
 

pkg/client/deploy/strategy.go

@@ -24,7 +24,9 @@ type Strategy interface {
 
 // RollingStrategy implements a rolling update deployment pattern where containers are updated one at a time
 // to minimize service disruption.
-type RollingStrategy struct{}
+type RollingStrategy struct {
+	State *scheduler.ClusterState
+}
 
 func (s *RollingStrategy) Type() string {
 	return "rolling"
@@ -33,12 +35,20 @@ func (s *RollingStrategy) Type() string {
 func (s *RollingStrategy) Plan(
 	ctx context.Context, cli scheduler.Client, svc *api.Service, spec api.ServiceSpec,
 ) (Plan, error) {
+	if s.State == nil {
+		state, err := scheduler.InspectClusterState(ctx, cli)
+		if err != nil {
+			return Plan{}, fmt.Errorf("inspect cluster state: %w", err)
+		}
+		s.State = state
+	}
+
 	// We can assume that the spec is valid at this point because it has been validated by the deployment.
 	switch spec.Mode {
 	case api.ServiceModeReplicated:
-		return s.planReplicated(ctx, cli, svc, spec)
+		return s.planReplicated(svc, spec)
 	case api.ServiceModeGlobal:
-		return s.planGlobal(ctx, cli, svc, spec)
+		return s.planGlobal(svc, spec)
 	default:
 		return Plan{}, fmt.Errorf("unsupported service mode: '%s'", spec.Mode)
 	}
@@ -48,18 +58,13 @@ func (s *RollingStrategy) Plan(
 // For replicated services, we want to maintain a specific number of containers (replicas) across the available machines
 // in the cluster.
 // TODO: schedule containers only on machines that contain the image if pull policy is set to 'never'.
-func (s *RollingStrategy) planReplicated(
-	ctx context.Context, cli scheduler.Client, svc *api.Service, spec api.ServiceSpec,
-) (Plan, error) {
+func (s *RollingStrategy) planReplicated(svc *api.Service, spec api.ServiceSpec) (Plan, error) {
 	plan, err := newEmptyPlan(svc, spec)
 	if err != nil {
 		return plan, err
 	}
 
-	sched, err := scheduler.NewServiceSchedulerWithClient(ctx, cli, spec)
-	if err != nil {
-		return plan, err
-	}
+	sched := scheduler.NewServiceScheduler(s.State, spec)
 	// TODO: return a detailed report on required constraints and which ones are satisfied?
 	availableMachines, err := sched.EligibleMachines()
 	if err != nil {
@@ -71,8 +76,6 @@ func (s *RollingStrategy) planReplicated(
 		matchedMachines = append(matchedMachines, m.Info)
 	}
 
-	// TODO: filter machines that contain the service volumes if the service uses any.
-
 	// Randomise the order of machines to avoid always deploying to the same machines first.
 	rand.Shuffle(len(matchedMachines), func(i, j int) {
 		matchedMachines[i], matchedMachines[j] = matchedMachines[j], matchedMachines[i]
@@ -198,9 +201,7 @@ func (s *RollingStrategy) planReplicated(
 // possible. If the new container would have port conflicts with the existing one, the old container is removed first.
 // It handles multiple containers per machine (though this should not occur in normal operation) and skips machines
 // that are down.
-func (s *RollingStrategy) planGlobal(
-	ctx context.Context, cli scheduler.Client, svc *api.Service, spec api.ServiceSpec,
-) (Plan, error) {
+func (s *RollingStrategy) planGlobal(svc *api.Service, spec api.ServiceSpec) (Plan, error) {
 	plan, err := newEmptyPlan(svc, spec)
 	if err != nil {
 		return plan, err
@@ -216,11 +217,7 @@ func (s *RollingStrategy) planGlobal(
 		}
 	}
 
-	sched, err := scheduler.NewServiceSchedulerWithClient(ctx, cli, spec)
-	if err != nil {
-		return plan, err
-	}
-
+	sched := scheduler.NewServiceScheduler(s.State, spec)
 	availableMachines, err := sched.EligibleMachines()
 	if err != nil {
 		return plan, err

pkg/client/service.go

@@ -38,7 +38,11 @@ func (cli *Client) RunService(ctx context.Context, spec api.ServiceSpec) (api.Ru
 
 	// Create missing named Docker volumes for the service.
 	if len(spec.MountedDockerVolumes()) > 0 {
-		volumeScheduler, err := scheduler.NewVolumeSchedulerWithClient(ctx, cli, []api.ServiceSpec{spec})
+		state, err := scheduler.InspectClusterState(ctx, cli)
+		if err != nil {
+			return resp, fmt.Errorf("inspect cluster state: %w", err)
+		}
+		volumeScheduler, err := scheduler.NewVolumeScheduler(state, []api.ServiceSpec{spec})
 		if err != nil {
 			return resp, fmt.Errorf("init volume scheduler: %w", err)
 		}