s3-csi/pkg/driver/nodeserver.go

package driver

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"

	"github.com/container-storage-interface/spec/lib/go/csi"
	"github.com/hanzoai/s3-csi/pkg/mountmanager"
	"github.com/hanzoai/s3/s3/glog"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"k8s.io/mount-utils"
)

// MounterFactory creates a Mounter for a volume. It is a field on NodeServer
// so tests can substitute a fake that does not touch the real mount service.
type MounterFactory func(volumeID string, readOnly bool, driver *S3Driver, volContext map[string]string) (Mounter, error)

// CapacityFn resolves the desired capacity (bytes) for a volume from the
// orchestrator. Tests can replace this to avoid real Kubernetes API calls.
type CapacityFn func(volumeID string) (int64, error)

// BindMountFn performs a bind mount of source onto target, optionally
// read-only. Used by Volume.Publish and overridden in tests.
type BindMountFn func(source, target string, readOnly bool) error

// HealthCheckFn reports whether a staging path has a live, responsive FUSE
// mount. Overridden in tests to simulate a crashed mount.
type HealthCheckFn func(stagingPath string) bool

type NodeServer struct {
	csi.UnimplementedNodeServer

	Driver *S3Driver

	// information about the managed volumes
	volumes       sync.Map
	volumeMutexes *KeyMutex

	// stopCh signals the health monitor goroutine to stop. Guarded by
	// stopOnce so NodeCleanup is safe to call from multiple shutdown paths.
	stopCh   chan struct{}
	stopOnce sync.Once

	// recoveryWg tracks in-flight per-volume recovery goroutines spawned
	// by checkAndRecoverVolumes. Tests can wait on it after driving a
	// sweep to synchronize on the recovery work.
	recoveryWg sync.WaitGroup

	// activeRecoveries holds the volumeIDs that currently have a
	// recovery goroutine in flight, used to prevent successive sweeps
	// from piling up new recovery goroutines on top of a hung one.
	activeRecoveries sync.Map // map[string]struct{}

	// Injectable factories / operations (overridden in tests).
	mounterFactory   MounterFactory
	capacityFn       CapacityFn
	isHealthyFn      HealthCheckFn
	cleanupStagingFn func(stagingPath string) error
	unmountFn        func(path string) error
	bindMountFn      BindMountFn
}

var _ = csi.NodeServer(&NodeServer{})

func (ns *NodeServer) NodeStageVolume(ctx context.Context, req *csi.NodeStageVolumeRequest) (*csi.NodeStageVolumeResponse, error) {
	volumeID := req.GetVolumeId()
	// mount the fs here
	stagingTargetPath := req.GetStagingTargetPath()

	glog.Infof("node stage volume %s to %s", volumeID, stagingTargetPath)

	// Check arguments
	if req.GetVolumeCapability() == nil {
		return nil, status.Error(codes.InvalidArgument, "Volume capability missing in request")
	}
	if !isValidVolumeCapabilities(ns.Driver.vcap, []*csi.VolumeCapability{req.GetVolumeCapability()}) {
		// return nil, status.Error(codes.InvalidArgument, "Volume capability not supported")
	}
	if volumeID == "" {
		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
	}
	if stagingTargetPath == "" {
		return nil, status.Error(codes.InvalidArgument, "Target path missing in request")
	}

	volumeMutex := ns.getVolumeMutex(volumeID)
	volumeMutex.Lock()
	defer volumeMutex.Unlock()

	// The volume has been staged and is in memory cache.
	if _, ok := ns.volumes.Load(volumeID); ok {
		glog.Infof("volume %s has been already staged", volumeID)
		return &csi.NodeStageVolumeResponse{}, nil
	}

	// Phase 2 Enhancement: Check if staging path exists but is stale/corrupted
	// This handles cases where:
	// 1. The CSI driver restarted and lost its in-memory state
	// 2. The FUSE process died leaving a stale mount
	if isStagingPathHealthy(stagingTargetPath) {
		// The staging path is healthy - rebuild the cache from the existing mount
		// This preserves the existing FUSE mount and avoids disrupting any published volumes
		glog.Infof("volume %s has existing healthy mount at %s, rebuilding cache", volumeID, stagingTargetPath)
		volume := ns.rebuildVolumeFromStaging(volumeID, stagingTargetPath)
		volume.volContext = req.GetVolumeContext()
		volume.readOnly = isVolumeReadOnly(req)
		ns.volumes.Store(volumeID, volume)
		glog.Infof("volume %s cache rebuilt from existing staging at %s", volumeID, stagingTargetPath)
		return &csi.NodeStageVolumeResponse{}, nil
	}

	// Check if there's a stale/corrupted mount that needs cleanup
	if _, err := os.Stat(stagingTargetPath); err == nil || mount.IsCorruptedMnt(err) {
		glog.Infof("volume %s has stale staging path at %s, cleaning up", volumeID, stagingTargetPath)
		if err := cleanupStaleStagingPath(stagingTargetPath); err != nil {
			ns.removeVolumeMutex(volumeID)
			return nil, status.Errorf(codes.Internal, "failed to cleanup stale staging path %s: %v", stagingTargetPath, err)
		}
	}

	volContext := req.GetVolumeContext()
	readOnly := isVolumeReadOnly(req)

	volume, err := ns.stageNewVolume(volumeID, stagingTargetPath, volContext, readOnly)
	if err != nil {
		// node stage is unsuccessful
		ns.removeVolumeMutex(volumeID)

		if os.IsPermission(err) {
			return nil, status.Error(codes.PermissionDenied, err.Error())
		}
		if strings.Contains(err.Error(), "invalid argument") {
			return nil, status.Error(codes.InvalidArgument, err.Error())
		}
		return nil, status.Error(codes.Internal, err.Error())
	}

	ns.volumes.Store(volumeID, volume)
	glog.Infof("volume %s successfully staged to %s", volumeID, stagingTargetPath)

	return &csi.NodeStageVolumeResponse{}, nil
}

func (ns *NodeServer) NodePublishVolume(ctx context.Context, req *csi.NodePublishVolumeRequest) (*csi.NodePublishVolumeResponse, error) {
	volumeID := req.GetVolumeId()
	targetPath := req.GetTargetPath()
	stagingTargetPath := req.GetStagingTargetPath()

	glog.Infof("node publish volume %s to %s", volumeID, targetPath)

	// Check arguments
	if req.GetVolumeCapability() == nil {
		return nil, status.Error(codes.InvalidArgument, "Volume capability missing in request")
	}
	if !isValidVolumeCapabilities(ns.Driver.vcap, []*csi.VolumeCapability{req.GetVolumeCapability()}) {
		// return nil, status.Error(codes.InvalidArgument, "Volume capability not supported")
	}
	if volumeID == "" {
		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
	}
	if targetPath == "" {
		return nil, status.Error(codes.InvalidArgument, "Target path missing in request")
	}
	if stagingTargetPath == "" {
		return nil, status.Error(codes.InvalidArgument, "Staging target path missing in request")
	}

	volumeMutex := ns.getVolumeMutex(volumeID)
	volumeMutex.Lock()
	defer volumeMutex.Unlock()

	volume, ok := ns.volumes.Load(volumeID)
	if !ok {
		// Phase 1: Self-healing for missing volume cache
		// This handles the case where the CSI driver restarted and lost its in-memory state,
		// but kubelet thinks the volume is already staged and directly calls NodePublishVolume.
		glog.Infof("volume %s not found in cache, attempting self-healing", volumeID)

		if isStagingPathHealthy(stagingTargetPath) {
			// The staging path is healthy - rebuild volume cache from existing mount
			glog.Infof("volume %s has healthy staging at %s, rebuilding cache", volumeID, stagingTargetPath)
			rebuiltVol := ns.rebuildVolumeFromStaging(volumeID, stagingTargetPath)
			rebuiltVol.volContext = req.GetVolumeContext()
			rebuiltVol.readOnly = isPublishVolumeReadOnly(req)
			volume = rebuiltVol
			ns.volumes.Store(volumeID, volume)
		} else {
			// The staging path is not healthy - we need to re-stage the volume
			// This requires volume context which we have from the request
			glog.Infof("volume %s staging path %s is not healthy, re-staging", volumeID, stagingTargetPath)

			// Clean up stale staging path if it exists
			if err := cleanupStaleStagingPath(stagingTargetPath); err != nil {
				ns.removeVolumeMutex(volumeID)
				return nil, status.Errorf(codes.Internal, "failed to cleanup stale staging path %s: %v", stagingTargetPath, err)
			}

			// Re-stage the volume using the shared helper
			volContext := req.GetVolumeContext()
			readOnly := isPublishVolumeReadOnly(req)

			newVolume, err := ns.stageNewVolume(volumeID, stagingTargetPath, volContext, readOnly)
			if err != nil {
				ns.removeVolumeMutex(volumeID)
				return nil, status.Errorf(codes.Internal, "failed to re-stage volume: %v", err)
			}

			ns.volumes.Store(volumeID, newVolume)
			volume = newVolume
			glog.Infof("volume %s successfully re-staged to %s", volumeID, stagingTargetPath)
		}
	}

	// When pod uses a volume in read-only mode, k8s will automatically
	// mount the volume as a read-only file system.
	vol := volume.(*Volume)
	if err := vol.Publish(stagingTargetPath, targetPath, req.GetReadonly()); err != nil {
		return nil, status.Error(codes.Internal, err.Error())
	}
	vol.AddPublishPath(targetPath, req.GetReadonly())

	glog.Infof("volume %s successfully published to %s", volumeID, targetPath)
	return &csi.NodePublishVolumeResponse{}, nil
}

// rebuildVolumeFromStaging creates a Volume struct from an existing healthy staging mount.
// This is used for self-healing when the CSI driver restarts but the FUSE mount is still active.
// Note: The returned Volume won't have an unmounter, so Unstage will need special handling.
func (ns *NodeServer) rebuildVolumeFromStaging(volumeID string, stagingPath string) *Volume {
	return &Volume{
		VolumeId:    volumeID,
		StagedPath:  stagingPath,
		driver:      ns.Driver,
		localSocket: mountmanager.LocalSocketPath(ns.Driver.volumeSocketDir, volumeID),
		bindMountFn: ns.bindMountFn,
		// mounter and unmounter are nil - this is intentional
		// The FUSE process is already running, we just need to track the volume
		// The mount service will have the mount tracked if it's still alive
	}
}

// isPublishVolumeReadOnly determines if a volume should be mounted read-only based on the publish request.
func isPublishVolumeReadOnly(req *csi.NodePublishVolumeRequest) bool {
	if req.GetReadonly() {
		return true
	}
	cap := req.GetVolumeCapability()
	if cap == nil || cap.GetAccessMode() == nil {
		return false
	}
	return isReadOnlyAccessMode(cap.GetAccessMode().Mode)
}

func (ns *NodeServer) NodeUnpublishVolume(ctx context.Context, req *csi.NodeUnpublishVolumeRequest) (*csi.NodeUnpublishVolumeResponse, error) {
	volumeID := req.GetVolumeId()
	targetPath := req.GetTargetPath()
	glog.Infof("node unpublish volume %s from %s", volumeID, targetPath)

	if volumeID == "" {
		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
	}

	if targetPath == "" {
		return nil, status.Error(codes.InvalidArgument, "Target path missing in request")
	}

	volumeMutex := ns.getVolumeMutex(volumeID)
	volumeMutex.Lock()
	defer volumeMutex.Unlock()

	volume, ok := ns.volumes.Load(volumeID)
	if !ok {
		glog.Warningf("volume %s hasn't been published", volumeID)

		// make sure there is no any garbage
		_ = mount.CleanupMountPoint(targetPath, mountutil, true)

		return &csi.NodeUnpublishVolumeResponse{}, nil
	}

	vol := volume.(*Volume)
	if err := vol.Unpublish(targetPath); err != nil {
		return nil, status.Error(codes.Internal, err.Error())
	}
	vol.RemovePublishPath(targetPath)

	glog.Infof("volume %s successfully unpublished from %s", volumeID, targetPath)

	return &csi.NodeUnpublishVolumeResponse{}, nil
}

func (ns *NodeServer) NodeGetInfo(ctx context.Context, req *csi.NodeGetInfoRequest) (*csi.NodeGetInfoResponse, error) {
	glog.V(3).Infof("node get info, node id: %s", ns.Driver.nodeID)

	return &csi.NodeGetInfoResponse{
		NodeId: ns.Driver.nodeID,
	}, nil
}

func (ns *NodeServer) NodeGetCapabilities(ctx context.Context, req *csi.NodeGetCapabilitiesRequest) (*csi.NodeGetCapabilitiesResponse, error) {
	glog.V(3).Infof("node get capabilities")

	return &csi.NodeGetCapabilitiesResponse{
		Capabilities: []*csi.NodeServiceCapability{
			{
				Type: &csi.NodeServiceCapability_Rpc{
					Rpc: &csi.NodeServiceCapability_RPC{
						Type: csi.NodeServiceCapability_RPC_STAGE_UNSTAGE_VOLUME,
					},
				},
			},
			{
				Type: &csi.NodeServiceCapability_Rpc{
					Rpc: &csi.NodeServiceCapability_RPC{
						Type: csi.NodeServiceCapability_RPC_EXPAND_VOLUME,
					},
				},
			},
		},
	}, nil
}

func (ns *NodeServer) NodeUnstageVolume(ctx context.Context, req *csi.NodeUnstageVolumeRequest) (*csi.NodeUnstageVolumeResponse, error) {
	volumeID := req.GetVolumeId()
	stagingTargetPath := req.GetStagingTargetPath()

	glog.Infof("node unstage volume %s from %s", volumeID, stagingTargetPath)

	// Check arguments
	if volumeID == "" {
		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
	}
	if stagingTargetPath == "" {
		return nil, status.Error(codes.InvalidArgument, "Target path missing in request")
	}

	volumeMutex := ns.getVolumeMutex(volumeID)
	volumeMutex.Lock()
	defer volumeMutex.Unlock()

	volume, ok := ns.volumes.Load(volumeID)
	if !ok {
		glog.Warningf("volume %s hasn't been staged", volumeID)

		// make sure there is no any garbage
		_ = mount.CleanupMountPoint(stagingTargetPath, mountutil, true)

		// Also clean up cache directory and socket if they exist
		CleanupVolumeResources(ns.Driver, volumeID)
	} else {
		if err := volume.(*Volume).Unstage(stagingTargetPath); err != nil {
			return nil, status.Error(codes.Internal, err.Error())
		} else {
			ns.volumes.Delete(volumeID)
		}
	}

	// remove mutex on successfull unstage
	ns.volumeMutexes.RemoveMutex(volumeID)

	glog.Infof("volume %s successfully unstaged from %s", volumeID, stagingTargetPath)

	return &csi.NodeUnstageVolumeResponse{}, nil
}

func (ns *NodeServer) NodeExpandVolume(ctx context.Context, req *csi.NodeExpandVolumeRequest) (*csi.NodeExpandVolumeResponse, error) {
	volumeID := req.GetVolumeId()
	volumePath := req.GetVolumePath()
	requiredBytes := req.GetCapacityRange().GetRequiredBytes()

	glog.Infof("node expand volume %s to %d bytes", req.GetVolumeId(), requiredBytes)

	if len(volumeID) == 0 {
		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
	}

	if len(volumePath) == 0 {
		return nil, status.Error(codes.InvalidArgument, "Volume path missing in request")
	}

	// TODO Check if volume exists
	// TODO Check if node exists

	volumeMutex := ns.getVolumeMutex(volumeID)
	volumeMutex.Lock()
	defer volumeMutex.Unlock()

	if volume, ok := ns.volumes.Load(volumeID); ok {
		if err := volume.(*Volume).Quota(requiredBytes); err != nil {
			return nil, err
		}
	}

	return &csi.NodeExpandVolumeResponse{}, nil
}

func (ns *NodeServer) NodeCleanup() {
	ns.stopOnce.Do(func() {
		close(ns.stopCh)
		// Wait for any in-flight recovery goroutines to finish so the
		// driver does not shut down mid-recovery, which could leave a
		// volume in a half-staged state.
		ns.recoveryWg.Wait()
		glog.Infof("node cleanup: health monitor stopped, mount service retains mounts across restarts")
	})
}

func (ns *NodeServer) getVolumeMutex(volumeID string) *sync.Mutex {
	return ns.volumeMutexes.GetMutex(volumeID)
}

func (ns *NodeServer) removeVolumeMutex(volumeID string) {
	ns.volumeMutexes.RemoveMutex(volumeID)
}

// stageNewVolume creates and stages a new volume with the given parameters.
// This is a helper method used by both NodeStageVolume and NodePublishVolume (for re-staging).
// Both the mounter and capacity lookup are resolved via NodeServer fields so
// tests can inject fakes that do not touch the real mount service or k8s API.
func (ns *NodeServer) stageNewVolume(volumeID, stagingTargetPath string, volContext map[string]string, readOnly bool) (*Volume, error) {
	effectiveVolContext := cloneVolumeContext(volContext)
	capacity, hasCapacity, err := ns.resolveVolumeCapacity(volumeID, effectiveVolContext)
	if err != nil {
		return nil, err
	}
	if hasCapacity {
		effectiveVolContext[volumeCapacityKey] = strconv.FormatInt(capacity, 10)
	}

	mounter, err := ns.mounterFactory(volumeID, readOnly, ns.Driver, effectiveVolContext)
	if err != nil {
		return nil, err
	}

	volume := NewVolume(volumeID, mounter, ns.Driver)
	volume.bindMountFn = ns.bindMountFn
	volume.volContext = effectiveVolContext
	volume.readOnly = readOnly
	if err := volume.Stage(stagingTargetPath); err != nil {
		return nil, err
	}

	// Apply quota if available.
	if hasCapacity {
		if err := volume.Quota(capacity); err != nil {
			glog.Warningf("failed to apply quota for volume %s: %v", volumeID, err)
			// Clean up the staged mount since we're returning an error
			if unstageErr := volume.Unstage(stagingTargetPath); unstageErr != nil {
				glog.Errorf("failed to unstage volume %s after quota failure: %v", volumeID, unstageErr)
			}
			return nil, err
		}
	} else {
		glog.V(4).Infof("volume capacity is unavailable for %s", volumeID)
	}

	return volume, nil
}

func cloneVolumeContext(volContext map[string]string) map[string]string {
	cloned := make(map[string]string, len(volContext)+1)
	for key, value := range volContext {
		cloned[key] = value
	}
	return cloned
}

func (ns *NodeServer) resolveVolumeCapacity(volumeID string, volContext map[string]string) (int64, bool, error) {
	if ns.capacityFn != nil {
		capacity, err := ns.capacityFn(volumeID)
		if err == nil {
			return capacity, true, nil
		}
		glog.V(4).Infof("could not resolve capacity for volume %s from orchestrator API, falling back to volume context: %v", volumeID, err)
	}

	value := volContext[volumeCapacityKey]
	if value == "" {
		return 0, false, nil
	}
	capacity, err := strconv.ParseInt(value, 10, 64)
	if err != nil || capacity < 0 {
		return 0, false, fmt.Errorf("invalid %s value %q", volumeCapacityKey, value)
	}
	return capacity, true, nil
}

// isReadOnlyAccessMode checks if the given access mode is read-only.
func isReadOnlyAccessMode(mode csi.VolumeCapability_AccessMode_Mode) bool {
	readOnlyModes := []csi.VolumeCapability_AccessMode_Mode{
		csi.VolumeCapability_AccessMode_SINGLE_NODE_READER_ONLY,
		csi.VolumeCapability_AccessMode_MULTI_NODE_READER_ONLY,
	}

	for _, readOnlyMode := range readOnlyModes {
		if mode == readOnlyMode {
			return true
		}
	}

	return false
}

func isVolumeReadOnly(req *csi.NodeStageVolumeRequest) bool {
	cap := req.GetVolumeCapability()
	if cap == nil || cap.GetAccessMode() == nil {
		return false
	}
	return isReadOnlyAccessMode(cap.GetAccessMode().Mode)
}