1 files changed, 59 insertions, 12 deletions

diff --git a/weed/operation/assign_file_id.go b/weed/operation/assign_file_id.go
index 2dfa44483..ffd3e4938 100644
--- a/weed/operation/assign_file_id.go
+++ b/weed/operation/assign_file_id.go
@@ -3,11 +3,14 @@ package operation
 import (
 	"context"
 	"fmt"
+	"strings"
+
+	"github.com/chrislusf/seaweedfs/weed/storage/needle"
+	"google.golang.org/grpc"
+
 	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
 	"github.com/chrislusf/seaweedfs/weed/security"
 	"github.com/chrislusf/seaweedfs/weed/util"
-	"google.golang.org/grpc"
-	"strings"
 )
 
 type VolumeAssignRequest struct {
@@ -15,6 +18,7 @@ type VolumeAssignRequest struct {
 	Replication         string
 	Collection          string
 	Ttl                 string
+	DiskType            string
 	DataCenter          string
 	Rack                string
 	DataNode            string
@@ -30,7 +34,7 @@ type AssignResult struct {
 	Auth      security.EncodedJwt `json:"auth,omitempty"`
 }
 
-func Assign(server string, grpcDialOption grpc.DialOption, primaryRequest *VolumeAssignRequest, alternativeRequests ...*VolumeAssignRequest) (*AssignResult, error) {
+func Assign(masterFn GetMasterFn, grpcDialOption grpc.DialOption, primaryRequest *VolumeAssignRequest, alternativeRequests ...*VolumeAssignRequest) (*AssignResult, error) {
 
 	var requests []*VolumeAssignRequest
 	requests = append(requests, primaryRequest)
@@ -44,17 +48,18 @@ func Assign(server string, grpcDialOption grpc.DialOption, primaryRequest *Volum
 			continue
 		}
 
-		lastError = WithMasterServerClient(server, grpcDialOption, func(masterClient master_pb.SeaweedClient) error {
+		lastError = WithMasterServerClient(masterFn(), grpcDialOption, func(masterClient master_pb.SeaweedClient) error {
 
 			req := &master_pb.AssignRequest{
-				Count:               primaryRequest.Count,
-				Replication:         primaryRequest.Replication,
-				Collection:          primaryRequest.Collection,
-				Ttl:                 primaryRequest.Ttl,
-				DataCenter:          primaryRequest.DataCenter,
-				Rack:                primaryRequest.Rack,
-				DataNode:            primaryRequest.DataNode,
-				WritableVolumeCount: primaryRequest.WritableVolumeCount,
+				Count:               request.Count,
+				Replication:         request.Replication,
+				Collection:          request.Collection,
+				Ttl:                 request.Ttl,
+				DiskType:            request.DiskType,
+				DataCenter:          request.DataCenter,
+				Rack:                request.Rack,
+				DataNode:            request.DataNode,
+				WritableVolumeCount: request.WritableVolumeCount,
 			}
 			resp, grpcErr := masterClient.Assign(context.Background(), req)
 			if grpcErr != nil {
@@ -81,6 +86,7 @@ func Assign(server string, grpcDialOption grpc.DialOption, primaryRequest *Volum
 			continue
 		}
 
+		break
 	}
 
 	return ret, lastError
@@ -99,3 +105,44 @@ func LookupJwt(master string, fileId string) security.EncodedJwt {
 
 	return security.EncodedJwt(tokenStr)
 }
+
+type StorageOption struct {
+	Replication       string
+	DiskType          string
+	Collection        string
+	DataCenter        string
+	Rack              string
+	TtlSeconds        int32
+	Fsync             bool
+	VolumeGrowthCount uint32
+}
+
+func (so *StorageOption) TtlString() string {
+	return needle.SecondsToTTL(so.TtlSeconds)
+}
+
+func (so *StorageOption) ToAssignRequests(count int) (ar *VolumeAssignRequest, altRequest *VolumeAssignRequest) {
+	ar = &VolumeAssignRequest{
+		Count:               uint64(count),
+		Replication:         so.Replication,
+		Collection:          so.Collection,
+		Ttl:                 so.TtlString(),
+		DiskType:            so.DiskType,
+		DataCenter:          so.DataCenter,
+		Rack:                so.Rack,
+		WritableVolumeCount: so.VolumeGrowthCount,
+	}
+	if so.DataCenter != "" || so.Rack != "" {
+		altRequest = &VolumeAssignRequest{
+			Count:               uint64(count),
+			Replication:         so.Replication,
+			Collection:          so.Collection,
+			Ttl:                 so.TtlString(),
+			DiskType:            so.DiskType,
+			DataCenter:          "",
+			Rack:                "",
+			WritableVolumeCount: so.VolumeGrowthCount,
+		}
+	}
+	return
+}