sambuc/feat merge restart pr 1 (#77)

* Merge Re-mount volumes on a restart or update #72

* Ensure the mutex is not copied, even when the nodeServer is copied by
  storing a pointer to the mutex, instead of the mutex itself

* Use Mutex instead of RWMutex, as having two readers of the variable at
  the same time means we are going to write the state at the same time,
  corrupting the state file on storage.

* Use private freestanding functions mutex as Mutex / RWMutex are not 
  recursive / re-entrant in Go.

* Improve the remounting logic by using a rate-limited amount of 
   background goroutines and adding timeouts values so that the service
   startup time is not dependent on the number of tracked volumes.

---------

Co-authored-by: Brian Kanya <[email protected]>

Author: sambuc

.devcontainer/rclone/install.sh

@@ -19,3 +19,6 @@ rm -rf /tmp/rclone
 # Fix the $GOPATH folder
 chown -R "${USERNAME}:golang" /go
 chmod -R g+r+w /go
+
+# Make sure the default folders exists
+mkdir -p /var/lib/kubelet/plugins/csi-rclone/

pkg/rclone/driver.go

@@ -1,9 +1,11 @@
 package rclone
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"os"
+	"path/filepath"
 	"sync"
 
 	"github.com/SwissDataScienceCenter/csi-rclone/pkg/kube"
@@ -78,14 +80,35 @@ func NewNodeServer(csiDriver *csicommon.CSIDriver, cacheDir string, cacheSize st
 	}
 	rcloneOps := NewRclone(kubeClient, rclonePort, cacheDir, cacheSize)
 
-	return &nodeServer{
+	// Use kubelet plugin directory for state persistence
+	stateFile := "/var/lib/kubelet/plugins/csi-rclone/mounted_volumes.json"
+
+	ns := &nodeServer{
 		DefaultNodeServer: csicommon.NewDefaultNodeServer(csiDriver),
 		mounter: &mount.SafeFormatAndMount{
 			Interface: mount.New(""),
 			Exec:      utilexec.New(),
 		},
-		RcloneOps: rcloneOps,
-	}, nil
+		RcloneOps:      rcloneOps,
+		mountedVolumes: make(map[string]MountedVolume),
+		mutex:          &sync.Mutex{},
+		stateFile:      stateFile,
+	}
+
+	// Ensure the folder exists
+	if err = os.MkdirAll(filepath.Dir(ns.stateFile), 0755); err != nil {
+		return nil, fmt.Errorf("failed to create state directory: %v", err)
+	}
+
+	// Load persisted state on startup
+	ns.mutex.Lock()
+	defer ns.mutex.Unlock()
+
+	if ns.mountedVolumes, err = readVolumeMap(ns.stateFile); err != nil {
+		klog.Warningf("Failed to load persisted volume state: %v", err)
+	}
+
+	return ns, nil
 }
 
 func NewControllerServer(csiDriver *csicommon.CSIDriver) *controllerServer {
@@ -139,7 +162,13 @@ func (d *Driver) Run() error {
 	)
 	d.server = s
 	if d.ns != nil && d.ns.RcloneOps != nil {
-		return d.ns.RcloneOps.Run()
+		onDaemonReady := func() error {
+			if d.ns != nil {
+				return d.ns.remountTrackedVolumes(context.Background())
+			}
+			return nil
+		}
+		return d.ns.RcloneOps.Run(onDaemonReady)
 	}
 	s.Wait()
 	return nil

pkg/rclone/nodeserver.go

@@ -7,10 +7,13 @@ package rclone
 
 import (
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
+	"runtime"
 	"strings"
+	"sync"
 	"time"
 
 	"gopkg.in/ini.v1"
@@ -37,6 +40,23 @@ type nodeServer struct {
 	*csicommon.DefaultNodeServer
 	mounter   *mount.SafeFormatAndMount
 	RcloneOps Operations
+
+	// Track mounted volumes for automatic remounting
+	mountedVolumes map[string]MountedVolume
+	mutex          *sync.Mutex
+	stateFile      string
+}
+
+type MountedVolume struct {
+	VolumeId        string
+	TargetPath      string
+	Remote          string
+	RemotePath      string
+	ConfigData      string
+	ReadOnly        bool
+	Parameters      map[string]string
+	SecretName      string
+	SecretNamespace string
 }
 
 // Mounting Volume (Preparation)
@@ -141,6 +161,10 @@ func (ns *nodeServer) NodePublishVolume(ctx context.Context, req *csi.NodePublis
 		}
 		return nil, status.Error(codes.Internal, err.Error())
 	}
+
+	// Track the mounted volume for automatic remounting
+	ns.trackMountedVolume(volumeId, targetPath, remote, remotePath, configData, readOnly, flags, secretName, secretNamespace)
+
 	// err = ns.WaitForMountAvailable(targetPath)
 	// if err != nil {
 	// 	return nil, status.Error(codes.Internal, err.Error())
@@ -323,6 +347,10 @@ func (ns *nodeServer) NodeUnpublishVolume(ctx context.Context, req *csi.NodeUnpu
 	if err := ns.RcloneOps.Unmount(ctx, req.GetVolumeId(), targetPath); err != nil {
 		klog.Warningf("Unmounting volume failed: %s", err)
 	}
+
+	// Remove the volume from tracking
+	ns.removeTrackedVolume(req.GetVolumeId())
+
 	mount.CleanupMountPoint(req.GetTargetPath(), ns.mounter, false)
 	return &csi.NodeUnpublishVolumeResponse{}, nil
 }
@@ -344,6 +372,122 @@ func (*nodeServer) NodeExpandVolume(ctx context.Context, req *csi.NodeExpandVolu
 	return nil, status.Errorf(codes.Unimplemented, "method NodeExpandVolume not implemented")
 }
 
+// Track mounted volume for automatic remounting
+func (ns *nodeServer) trackMountedVolume(volumeId, targetPath, remote, remotePath, configData string, readOnly bool, parameters map[string]string, secretName, secretNamespace string) {
+	ns.mutex.Lock()
+	defer ns.mutex.Unlock()
+
+	ns.mountedVolumes[volumeId] = MountedVolume{
+		VolumeId:        volumeId,
+		TargetPath:      targetPath,
+		Remote:          remote,
+		RemotePath:      remotePath,
+		ConfigData:      configData,
+		ReadOnly:        readOnly,
+		Parameters:      parameters,
+		SecretName:      secretName,
+		SecretNamespace: secretNamespace,
+	}
+	klog.Infof("Tracked mounted volume %s at path %s", volumeId, targetPath)
+
+	if err := writeVolumeMap(ns.stateFile, ns.mountedVolumes); err != nil {
+		klog.Errorf("Failed to persist volume state: %v", err)
+	}
+}
+
+// Remove tracked volume when unmounted
+func (ns *nodeServer) removeTrackedVolume(volumeId string) {
+	ns.mutex.Lock()
+	defer ns.mutex.Unlock()
+
+	delete(ns.mountedVolumes, volumeId)
+	klog.Infof("Removed tracked volume %s", volumeId)
+
+	if err := writeVolumeMap(ns.stateFile, ns.mountedVolumes); err != nil {
+		klog.Errorf("Failed to persist volume state: %v", err)
+	}
+}
+
+// Automatically remount all tracked volumes after daemon restart
+func (ns *nodeServer) remountTrackedVolumes(ctx context.Context) error {
+	type mountResult struct {
+		volumeID string
+		err      error
+	}
+
+	ns.mutex.Lock()
+	defer ns.mutex.Unlock()
+
+	if len(ns.mountedVolumes) == 0 {
+		klog.Info("No tracked volumes to remount")
+		return nil
+	}
+
+	klog.Infof("Remounting %d tracked volumes", len(ns.mountedVolumes))
+
+	// Limit the number of active workers to the number of CPU threads (arbitrarily chosen)
+	limits := make(chan bool, runtime.GOMAXPROCS(0))
+	defer close(limits)
+
+	volumesCount := len(ns.mountedVolumes)
+	results := make(chan mountResult, volumesCount)
+	defer close(results)
+
+	ctxWithTimeout, cancel := context.WithTimeout(ctx, 60*time.Second)
+	defer cancel()
+
+	for volumeId, mv := range ns.mountedVolumes {
+		go func() {
+			limits <- true // block until there is a free slot in the queue
+			defer func() {
+				<-limits // free a slot in the queue when we exit
+			}()
+
+			ctxWithMountTimeout, cancel := context.WithTimeout(ctxWithTimeout, 30*time.Second)
+			defer cancel()
+
+			klog.Infof("Remounting volume %s to %s", volumeId, mv.TargetPath)
+
+			// Create the mount directory if it doesn't exist
+			var err error
+			if err = os.MkdirAll(mv.TargetPath, 0750); err != nil {
+				klog.Errorf("Failed to create mount directory %s: %v", mv.TargetPath, err)
+			} else {
+				// Remount the volume
+				rcloneVol := &RcloneVolume{
+					ID:         mv.VolumeId,
+					Remote:     mv.Remote,
+					RemotePath: mv.RemotePath,
+				}
+
+				err = ns.RcloneOps.Mount(ctxWithMountTimeout, rcloneVol, mv.TargetPath, mv.ConfigData, mv.ReadOnly, mv.Parameters)
+			}
+
+			results <- mountResult{volumeId, err}
+		}()
+	}
+
+	for {
+		select {
+		case result := <-results:
+			volumesCount--
+			if result.err != nil {
+				klog.Errorf("Failed to remount volume %s: %v", result.volumeID, result.err)
+				// Don't return error here, continue with other volumes not to block all users because of a failed mount.
+				delete(ns.mountedVolumes, result.volumeID)
+				// Should we keep it on disk? This will be lost on the first new mount which will override the file.
+			} else {
+				klog.Infof("Successfully remounted volume %s", result.volumeID)
+			}
+			if volumesCount == 0 {
+				return nil
+			}
+		case <-ctxWithTimeout.Done():
+			return ctxWithTimeout.Err()
+		}
+	}
+}
+
 func (ns *nodeServer) WaitForMountAvailable(mountpoint string) error {
 	for {
 		select {
@@ -357,3 +501,47 @@ func (ns *nodeServer) WaitForMountAvailable(mountpoint string) error {
 		}
 	}
 }
+
+// Persist volume state to disk
+func writeVolumeMap(filename string, volumes map[string]MountedVolume) error {
+	if filename == "" {
+		return nil
+	}
+
+	data, err := json.Marshal(volumes)
+	if err != nil {
+		return fmt.Errorf("failed to marshal volume state: %v", err)
+	}
+
+	if err := os.WriteFile(filename, data, 0600); err != nil {
+		return fmt.Errorf("failed to write state file: %v", err)
+	}
+
+	klog.Infof("Persisted volume state to %s", filename)
+	return nil
+}
+
+// Load volume state from disk
+func readVolumeMap(filename string) (map[string]MountedVolume, error) {
+	volumes := make(map[string]MountedVolume)
+
+	if filename == "" {
+		return volumes, nil
+	}
+
+	data, err := os.ReadFile(filename)
+	if err != nil {
+		if os.IsNotExist(err) {
+			klog.Info("No persisted volume state found, starting fresh")
+			return volumes, nil
+		}
+		return volumes, fmt.Errorf("failed to read state file: %v", err)
+	}
+
+	if err := json.Unmarshal(data, &volumes); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal volume state: %v", err)
+	}
+
+	klog.Infof("Loaded %d tracked volumes from %s", len(volumes), filename)
+	return volumes, nil
+}

pkg/rclone/rclone.go

@@ -34,7 +34,7 @@ type Operations interface {
 	Unmount(ctx context.Context, volumeId string, targetPath string) error
 	GetVolumeById(ctx context.Context, volumeId string) (*RcloneVolume, error)
 	Cleanup() error
-	Run() error
+	Run(onDaemonReady func() error) error
 }
 
 type Rclone struct {
@@ -472,11 +472,16 @@ func (r *Rclone) start_daemon() error {
 	return nil
 }
 
-func (r *Rclone) Run() error {
+func (r *Rclone) Run(onDaemonReady func() error) error {
 	err := r.start_daemon()
 	if err != nil {
 		return err
 	}
+	if onDaemonReady != nil {
+		if err := onDaemonReady(); err != nil {
+			klog.Warningf("Error in onDaemonReady callback: %v", err)
+		}
+	}
 	// blocks until the rclone daemon is stopped
 	return r.daemonCmd.Wait()
 }