Add dataLocality flag, checking and handling

3 files changed, 63 insertions, 0 deletions

diff --git a/cmd/seaweedfs-csi-driver/main.go b/cmd/seaweedfs-csi-driver/main.go
index 99189d0..3a0f2c8 100644
--- a/cmd/seaweedfs-csi-driver/main.go
+++ b/cmd/seaweedfs-csi-driver/main.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"os"
 
+	"github.com/seaweedfs/seaweedfs-csi-driver/pkg/datalocality"
 	"github.com/seaweedfs/seaweedfs-csi-driver/pkg/driver"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	flag "github.com/seaweedfs/seaweedfs/weed/util/fla9"
@@ -20,6 +21,9 @@ var (
 	cacheDir          = flag.String("cacheDir", os.TempDir(), "local cache directory for file chunks and meta data")
 	uidMap            = flag.String("map.uid", "", "map local uid to uid on filer, comma-separated <local_uid>:<filer_uid>")
 	gidMap            = flag.String("map.gid", "", "map local gid to gid on filer, comma-separated <local_gid>:<filer_gid>")
+	dataCenter		  = flag.String("dataCenter", "", "dataCenter this node is running in (locality-definition)")
+	dataLocalityStr	  = flag.String("dataLocality", "", "which volume-nodes pods will use for activity (one-of: 'write_preferLocalDc'). Requires used locality-definitions to be set")
+	dataLocality      datalocality.DataLocality
 )
 
 func main() {
@@ -35,6 +39,18 @@ func main() {
 		os.Exit(0)
 	}
 
+	err := convertRequiredValues()
+	if(err != nil){
+		glog.Error("Failed converting flag: ", err);
+		os.Exit(1);
+	}
+
+	err = checkPreconditions()
+	if(err != nil){
+		glog.Error("Precondition failed: ", err);
+		os.Exit(1);
+	}
+
 	glog.Infof("connect to filer %s", *filer)
 
 	drv := driver.NewSeaweedFsDriver(*filer, *nodeID, *endpoint)
@@ -43,5 +59,29 @@ func main() {
 	drv.CacheDir = *cacheDir
 	drv.UidMap = *uidMap
 	drv.GidMap = *gidMap
+	drv.DataCenter = *dataCenter
+	drv.DataLocality = dataLocality
+
 	drv.Run()
 }
+
+func convertRequiredValues() error {
+	// Convert DataLocalityStr to DataLocality
+	if(*dataLocalityStr != ""){
+		var ok bool
+		dataLocality, ok = datalocality.FromString(*dataLocalityStr)
+		if(!ok){
+			return fmt.Errorf("dataLocality invalid value")
+		}
+	}
+
+	return nil
+}
+
+func checkPreconditions() error {
+	if(dataLocality != datalocality.None && *dataCenter == ""){
+		return fmt.Errorf("dataLocality set, but not all locality-definitions were set! ('dataCenter')")
+	}
+
+	return nil
+}
diff --git a/pkg/driver/driver.go b/pkg/driver/driver.go
index 47a55e0..1cc3832 100644
--- a/pkg/driver/driver.go
+++ b/pkg/driver/driver.go
@@ -10,6 +10,7 @@ import (
 	"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb"
 	"github.com/seaweedfs/seaweedfs/weed/security"
 	"github.com/seaweedfs/seaweedfs/weed/util"
+	"github.com/seaweedfs/seaweedfs-csi-driver/pkg/datalocality"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -43,6 +44,8 @@ type SeaweedFsDriver struct {
 	CacheDir          string
 	UidMap            string
 	GidMap            string
+	DataCenter        string
+	DataLocality	  datalocality.DataLocality
 }
 
 func NewSeaweedFsDriver(filer, nodeID, endpoint string) *SeaweedFsDriver {
diff --git a/pkg/driver/mounter_seaweedfs.go b/pkg/driver/mounter_seaweedfs.go
index 7035cd4..cbb77cd 100644
--- a/pkg/driver/mounter_seaweedfs.go
+++ b/pkg/driver/mounter_seaweedfs.go
@@ -7,6 +7,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/seaweedfs/seaweedfs-csi-driver/pkg/datalocality"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"github.com/seaweedfs/seaweedfs/weed/util"
 )
@@ -85,6 +86,24 @@ func (seaweedFs *seaweedFsMounter) Mount(target string) (Unmounter, error) {
 		"map.gid":				seaweedFs.driver.GidMap,
 	}
 
+	// Handle DataLocality
+	dataLocality := seaweedFs.driver.DataLocality;
+	// Try to override when set in context
+	if dataLocalityStr, ok := seaweedFs.volContext["dataLocality"]; ok{
+		// Convert to enum
+		dataLocalityRes, ok := datalocality.FromString(dataLocalityStr)
+		if(!ok){
+			glog.Warning("volumeContext 'dataLocality' invalid");
+		}else{
+			dataLocality = dataLocalityRes
+		}
+	}
+	// Settings based on type
+	switch(dataLocality){
+	case datalocality.Write_preferLocalDc:
+		argsMap["dataCenter"] = seaweedFs.driver.DataCenter;
+	}
+
 	// volContext-parameter -> mount-arg
 	parameterArgMap := map[string]string{
 		"uidMap":		"map.uid",
@@ -97,6 +116,7 @@ func (seaweedFs *seaweedFsMounter) Mount(target string) (Unmounter, error) {
 	// Fields supplied in context, but ignored because they are handled explicitly somewhere else
 	ignoreArgs := []string{
 		"volumeCapacity",
+		"dataLocality",
 	}
 
 	//	Merge volContext into argsMap with key-mapping