Merge pull request #102 from Ruakij/feat_dataLocality

4 files changed, 66 insertions, 1 deletions

diff --git a/pkg/datalocality/mapping.go b/pkg/datalocality/mapping.go
new file mode 100644
index 0000000..51484b0
--- /dev/null
+++ b/pkg/datalocality/mapping.go
@@ -0,0 +1,31 @@
+package datalocality
+
+import (
+	"strings"
+)
+
+type DataLocality uint
+
+const (
+	None				DataLocality = iota
+	Write_preferLocalDc
+)
+
+// DataLocality -> String
+var dataLocalityStringMap = []string {
+	"none",
+	"write_preferlocaldc",
+}
+func (d DataLocality) String() string {
+	return dataLocalityStringMap[d]
+}
+
+// String -> DataLocality
+var stringDataLocalityMap = map[string]DataLocality {
+	"none": None,
+	"write_preferlocaldc": Write_preferLocalDc,
+}
+func FromString(s string) (DataLocality, bool) {
+	value, ok := stringDataLocalityMap[strings.ToLower(s)]
+	return value, ok
+}
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 16d97c8..12142af 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"
 )
@@ -93,6 +94,27 @@ func (seaweedFs *seaweedFsMounter) Mount(target string) (Unmounter, error) {
 		"readRetryTime":		"",
 	}
 
+	// 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
+		}
+	}
+	if err := CheckDataLocality(&dataLocality, &seaweedFs.driver.DataCenter); err != nil {
+		return nil, err
+	}
+	// 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",
@@ -101,10 +123,11 @@ func (seaweedFs *seaweedFsMounter) Mount(target string) (Unmounter, error) {
 		// volumeContext has "diskType", but mount-option is "disk", converting for backwards compatability
 		"diskType":		"disk",
 	}
-	
+
 	// Explicitly ignored volContext args e.g. handled somewhere else
 	ignoreArgs := []string{
 		"volumeCapacity",
+		"dataLocality",
 	}
 
 	//	Merge volContext into argsMap with key-mapping
diff --git a/pkg/driver/utils.go b/pkg/driver/utils.go
index b629543..e103079 100644
--- a/pkg/driver/utils.go
+++ b/pkg/driver/utils.go
@@ -8,6 +8,7 @@ import (
 	"sync"
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
+	"github.com/seaweedfs/seaweedfs-csi-driver/pkg/datalocality"
 	"github.com/seaweedfs/seaweedfs/weed/glog"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
@@ -121,3 +122,10 @@ func (km *KeyMutex) GetMutex(key string) *sync.Mutex {
 func (km *KeyMutex) RemoveMutex(key string) {
 	km.mutexes.Delete(key)
 }
+
+func CheckDataLocality(dataLocality *datalocality.DataLocality, dataCenter *string) error {
+	if(*dataLocality != datalocality.None && *dataCenter == ""){
+		return fmt.Errorf("dataLocality set, but not all locality-definitions were set")
+	}
+	return nil
+}