CSI volume support: follow-up improvements

cmd/ateapi/internal/controlapi/functionaltest/actor_test.go

@@ -1492,8 +1492,8 @@ func (f *partialFailVolumePlugin) CreateVolume(ctx context.Context, name, capaci
 	return "storage-" + name, parameters, nil
 }
 
-func (f *partialFailVolumePlugin) AttachVolume(ctx context.Context, volumeID, node string) error {
-	return nil
+func (f *partialFailVolumePlugin) AttachVolume(ctx context.Context, volumeID, node string, mode ateapipb.VolumeAccessMode) (map[string]string, error) {
+	return nil, nil
 }
 
 func (f *partialFailVolumePlugin) DetachVolume(ctx context.Context, volumeID, node string) error {
@@ -1633,8 +1633,8 @@ func (r *retrySuccessVolumePlugin) CreateVolume(ctx context.Context, name, capac
 	return "storage-" + name, parameters, nil
 }
 
-func (r *retrySuccessVolumePlugin) AttachVolume(ctx context.Context, volumeID, node string) error {
-	return nil
+func (r *retrySuccessVolumePlugin) AttachVolume(ctx context.Context, volumeID, node string, mode ateapipb.VolumeAccessMode) (map[string]string, error) {
+	return nil, nil
 }
 
 func (r *retrySuccessVolumePlugin) DetachVolume(ctx context.Context, volumeID, node string) error {

cmd/ateapi/internal/controlapi/volumes.go

@@ -21,6 +21,7 @@ import (
 	"log/slog"
 
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
+	"github.com/agent-substrate/substrate/internal/volume"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -49,6 +50,7 @@ func initialActorVolumes(ctx context.Context, scLister storagev1listers.StorageC
 				VolumeName: vol.GetName(),
 				VolumeType: sc.Provisioner,
 				Status:     ateapipb.ExternalVolume_STATUS_PENDING,
+				AccessMode: vol.GetExternalVolumeTemplate().GetAccessMode(),
 			})
 		}
 	}
@@ -109,7 +111,7 @@ func createActorVolumes(ctx context.Context, registry VolumePluginRegistry, scLi
 			return resultVolumes, status.Errorf(codes.FailedPrecondition, "volume %q has mismatched type %q (expected %q from StorageClass %q)", volName, vol.GetVolumeType(), sc.Provisioner, scName)
 		}
 
-		plugin, err := registry.GetPlugin(ctx, vol.GetVolumeType())
+		plugin, err := volume.LookupPlugin(ctx, registry.GetPlugin, vol.GetVolumeType())
 		if err != nil {
 			return resultVolumes, status.Errorf(codes.FailedPrecondition, "failed to get volume plugin for driver %q (StorageClass %q): %v", sc.Provisioner, scName, err)
 		}
@@ -125,6 +127,7 @@ func createActorVolumes(ctx context.Context, registry VolumePluginRegistry, scLi
 			VolumeType:      sc.Provisioner,
 			Status:          ateapipb.ExternalVolume_STATUS_CREATED,
 			VolumeContext:   volCtx,
+			AccessMode:      specVol.GetExternalVolumeTemplate().GetAccessMode(),
 		})
 	}
 	return resultVolumes, nil
@@ -144,11 +147,9 @@ func deleteActorVolumes(ctx context.Context, registry VolumePluginRegistry, acto
 			// to the original requested volID.
 			volID = actorVolumeID(actorUID, vol.GetVolumeName())
 		}
-		// TODO: Standardize volume plugin lookup and error handling across control plane
-		// and worker plane (e.g. via a shared helper).
-		plugin, err := registry.GetPlugin(ctx, vol.GetVolumeType())
+		plugin, err := volume.LookupPlugin(ctx, registry.GetPlugin, vol.GetVolumeType())
 		if err != nil {
-			errs = append(errs, fmt.Errorf("failed to get volume plugin for %q: %w", vol.GetVolumeType(), err))
+			errs = append(errs, err)
 			continue
 		}
 		if err := plugin.DeleteVolume(ctx, volID); err != nil {
@@ -235,17 +236,22 @@ func detachActorVolumes(ctx context.Context, st detachActorVolumesStore, registr
 			continue
 		}
 		slog.InfoContext(ctx, "Detaching volume from node", slog.String("volume_id", vol.GetStorageVolumeId()), slog.String("node", node))
-		plugin, err := registry.GetPlugin(ctx, vol.GetVolumeType())
+		plugin, err := volume.LookupPlugin(ctx, registry.GetPlugin, vol.GetVolumeType())
 		if err != nil {
-			errs = append(errs, fmt.Errorf("failed to get volume plugin for %q: %w", vol.GetVolumeType(), err))
+			errs = append(errs, err)
 			continue
 		}
 		if err := plugin.DetachVolume(ctx, vol.GetStorageVolumeId(), node); err != nil {
 			if status.Code(err) == codes.NotFound {
 				slog.WarnContext(ctx, "Volume not found during detach, assuming already detached", slog.String("volume_id", vol.GetStorageVolumeId()), slog.String("node", node))
+				vol.PublishContext = nil
+				vol.PublishContextNode = ""
 				continue
 			}
 			errs = append(errs, fmt.Errorf("failed to detach volume %q from node %q: %w", vol.GetStorageVolumeId(), node, err))
+		} else {
+			vol.PublishContext = nil
+			vol.PublishContextNode = ""
 		}
 	}
 	return errors.Join(errs...)

cmd/ateapi/internal/controlapi/volumes_test.go

@@ -56,6 +56,7 @@ func TestInitialActorVolumes_PendingState(t *testing.T) {
 				Name: "data-vol-1",
 				ExternalVolumeTemplate: &ateapipb.ExternalVolumeTemplate{
 					StorageClassName: "standard",
+					AccessMode:       ateapipb.VolumeAccessMode_VOLUME_ACCESS_MODE_READ_ONLY_MANY,
 				},
 			},
 			{
@@ -79,6 +80,7 @@ func TestInitialActorVolumes_PendingState(t *testing.T) {
 			VolumeName: "data-vol-1",
 			VolumeType: "mock-standard",
 			Status:     ateapipb.ExternalVolume_STATUS_PENDING,
+			AccessMode: ateapipb.VolumeAccessMode_VOLUME_ACCESS_MODE_READ_ONLY_MANY,
 		},
 		{
 			VolumeName: "data-vol-2",

cmd/ateapi/internal/controlapi/workflow_resume.go

@@ -26,6 +26,7 @@ import (
 	"github.com/agent-substrate/substrate/internal/ateattr"
 	"github.com/agent-substrate/substrate/internal/proto/ateletpb"
 	"github.com/agent-substrate/substrate/internal/resources"
+	"github.com/agent-substrate/substrate/internal/volume"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -121,9 +122,11 @@ func (w *ActorWorkflow) ResumeActor(ctx context.Context, actorRef resources.Acto
 		return nil, false, err
 	}
 	actor = assigned
-	if err = w.ensureVolumesAttached(leaseCtx, actor, worker, actorTemplate); err != nil {
+	var attached *ateapipb.Actor
+	if attached, err = w.ensureVolumesAttached(leaseCtx, actorRef, actor, worker, actorTemplate); err != nil {
 		return nil, false, err
 	}
+	actor = attached
 	if tele, err = w.ensureAteletRestored(leaseCtx, actorRef, actor, actorTemplate, src); err != nil {
 		return nil, false, err
 	}
@@ -622,27 +625,62 @@ func schedulingConstraints(actor *ateapipb.Actor, tmpl *ateapipb.ActorTemplate)
 // assigned worker's node. Attachment is idempotent, so a re-entered workflow
 // safely runs it again.
 // TODO replace re-execution with a proper check on the volumes' attach state.
-func (w *ActorWorkflow) ensureVolumesAttached(ctx context.Context, actor *ateapipb.Actor, worker *ateapipb.Worker, actorTemplate *ateapipb.ActorTemplate) (err error) {
+func (w *ActorWorkflow) ensureVolumesAttached(ctx context.Context, actorRef resources.ActorRef, actor *ateapipb.Actor, worker *ateapipb.Worker, actorTemplate *ateapipb.ActorTemplate) (_ *ateapipb.Actor, err error) {
 	ctx, done := stepSpan(ctx, "AttachVolumes")
 	defer func() { err = done(err) }()
 
 	node := worker.GetNodeName()
 	if node == "" {
-		return fmt.Errorf("assigned worker has no node name")
+		return nil, fmt.Errorf("assigned worker has no node name")
 	}
 
 	ref := &ateapipb.ObjectRef{Atespace: actor.GetMetadata().GetAtespace(), Name: actor.GetMetadata().GetName()}
-	for _, vol := range getMountedActorVolumes(ctx, ref, actor.GetStatus().GetActorVolumes(), actorTemplate) {
+	mountedVols := getMountedActorVolumes(ctx, ref, actor.GetStatus().GetActorVolumes(), actorTemplate)
+	if len(mountedVols) == 0 {
+		return actor, nil
+	}
+
+	updated := false
+	for _, vol := range mountedVols {
 		slog.InfoContext(ctx, "Attaching volume to node", slog.String("volume_id", vol.GetStorageVolumeId()), slog.String("node", node))
-		plugin, err := w.pluginRegistry.GetPlugin(ctx, vol.GetVolumeType())
+		plugin, err := volume.LookupPlugin(ctx, w.pluginRegistry.GetPlugin, vol.GetVolumeType())
 		if err != nil {
-			return fmt.Errorf("failed to get volume plugin for %q: %w", vol.GetVolumeType(), err)
+			return nil, err
 		}
-		if err := plugin.AttachVolume(ctx, vol.GetStorageVolumeId(), node); err != nil {
-			return fmt.Errorf("failed to attach volume %q to node %q: %w", vol.GetStorageVolumeId(), node, err)
+		pubCtx, err := plugin.AttachVolume(ctx, vol.GetStorageVolumeId(), node, vol.GetAccessMode())
+		if err != nil {
+			return nil, fmt.Errorf("failed to attach volume %q to node %q: %w", vol.GetStorageVolumeId(), node, err)
+		}
+		if vol.GetPublishContextNode() != node || len(pubCtx) > 0 {
+			vol.PublishContext = pubCtx
+			vol.PublishContextNode = node
+			updated = true
 		}
 	}
-	return nil
+
+	if updated {
+		updatePrecondition := store.PreconditionFrom(actor)
+		storedActor, updateErr := w.store.UpdateActor(ctx, actorRef, updatePrecondition, func(toUpdate *ateapipb.Actor) error {
+			for _, mVol := range mountedVols {
+				for _, toVol := range toUpdate.GetStatus().GetActorVolumes() {
+					if toVol.GetVolumeName() == mVol.GetVolumeName() {
+						toVol.PublishContext = mVol.GetPublishContext()
+						toVol.PublishContextNode = mVol.GetPublishContextNode()
+					}
+				}
+			}
+			return nil
+		})
+		if updateErr != nil {
+			if errors.Is(updateErr, store.ErrVersionConflict) {
+				return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
+			}
+			return nil, fmt.Errorf("while updating actor after volume attachment: %w", updateErr)
+		}
+		return storedActor, nil
+	}
+
+	return actor, nil
 }
 
 // ensureAteletRestored brings the workload up on the assigned worker:

cmd/ateapi/internal/controlapi/workload_spec.go

@@ -173,24 +173,32 @@ func appendExternalVolumes(workloadSpec *ateletpb.WorkloadSpec, template *ateapi
 			var storageVolID string
 			var volType string
 			var volCtx map[string]string
+			var pubCtx map[string]string
+			var accessMode ateapipb.VolumeAccessMode
 			for _, dbVol := range actor.GetStatus().GetActorVolumes() {
 				if dbVol.GetVolumeName() == vol.GetName() {
 					storageVolID = dbVol.GetStorageVolumeId()
 					volType = dbVol.GetVolumeType()
 					volCtx = dbVol.GetVolumeContext()
+					pubCtx = dbVol.GetPublishContext()
+					accessMode = dbVol.GetAccessMode()
 					break
 				}
 			}
 			if storageVolID == "" {
 				return fmt.Errorf("volume %s not found for actor %s", vol.GetName(), actor.GetMetadata().GetName())
 			}
+			isReadOnly := accessMode == ateapipb.VolumeAccessMode_VOLUME_ACCESS_MODE_READ_ONLY_MANY
 			workloadSpec.Volumes = append(workloadSpec.Volumes, &ateletpb.Volume{
 				Name: vol.GetName(),
 				Source: &ateletpb.Volume_External{
 					External: &ateletpb.ExternalVolumeSource{
 						StorageVolumeId: storageVolID,
 						VolumeType:      volType,
 						VolumeContext:   volCtx,
+						PublishContext:  pubCtx,
+						AccessMode:      accessMode,
+						Readonly:        isReadOnly,
 					},
 				},
 			})

cmd/ateapi/internal/controlapi/workload_spec_test.go

@@ -427,6 +427,8 @@ func TestAppendExternalVolumes(t *testing.T) {
 					StorageVolumeId: "vol-gce-pd-123",
 					VolumeType:      "pd-standard",
 					VolumeContext:   map[string]string{"foo": "bar"},
+					PublishContext:  map[string]string{"devicePath": "/dev/sda"},
+					AccessMode:      ateapipb.VolumeAccessMode_VOLUME_ACCESS_MODE_READ_ONLY_MANY,
 				},
 			},
 		},
@@ -446,6 +448,9 @@ func TestAppendExternalVolumes(t *testing.T) {
 						StorageVolumeId: "vol-gce-pd-123",
 						VolumeType:      "pd-standard",
 						VolumeContext:   map[string]string{"foo": "bar"},
+						PublishContext:  map[string]string{"devicePath": "/dev/sda"},
+						AccessMode:      ateapipb.VolumeAccessMode_VOLUME_ACCESS_MODE_READ_ONLY_MANY,
+						Readonly:        true,
 					},
 				},
 			},