From 61fa485700f15b537611dfd399890a63e4b61601 Mon Sep 17 00:00:00 2001
From: zhangsong <songz.qqmail@qq.com>
Date: Sun, 10 Nov 2019 20:11:03 +0800
Subject: [PATCH] add volume number param in assign operation

---
 weed/operation/assign_file_id.go         | 16 +++++++++-------
 weed/pb/master.proto                     |  1 +
 weed/pb/master_pb/master.pb.go           | 17 +++++++++--------
 weed/server/master_grpc_server_volume.go |  2 +-
 weed/server/master_server_handlers.go    |  7 ++++++-
 weed/topology/volume_growth.go           |  7 +++++--
 6 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/weed/operation/assign_file_id.go b/weed/operation/assign_file_id.go
index 4c50eaa26..2971cddbb 100644
--- a/weed/operation/assign_file_id.go
+++ b/weed/operation/assign_file_id.go
@@ -11,13 +11,14 @@ import (
 )
 
 type VolumeAssignRequest struct {
-	Count       uint64
-	Replication string
-	Collection  string
-	Ttl         string
-	DataCenter  string
-	Rack        string
-	DataNode    string
+	Count               uint64
+	Replication         string
+	Collection          string
+	Ttl                 string
+	DataCenter          string
+	Rack                string
+	DataNode            string
+	WritableVolumeCount uint32
 }
 
 type AssignResult struct {
@@ -53,6 +54,7 @@ func Assign(server string, grpcDialOption grpc.DialOption, primaryRequest *Volum
 				DataCenter:  primaryRequest.DataCenter,
 				Rack:        primaryRequest.Rack,
 				DataNode:    primaryRequest.DataNode,
+				WritableVolumeCount: primaryRequest.WritableVolumeCount,
 			}
 			resp, grpcErr := masterClient.Assign(context.Background(), req)
 			if grpcErr != nil {
diff --git a/weed/pb/master.proto b/weed/pb/master.proto
index 80231a2d8..0286bca49 100644
--- a/weed/pb/master.proto
+++ b/weed/pb/master.proto
@@ -140,6 +140,7 @@ message AssignRequest {
     string rack = 6;
     string data_node = 7;
     uint32 memory_map_max_size_mb = 8;
+    uint32 WritableVolumeCount = 9;
 }
 message AssignResponse {
     string fid = 1;
diff --git a/weed/pb/master_pb/master.pb.go b/weed/pb/master_pb/master.pb.go
index 880d39de5..8997f3a99 100644
--- a/weed/pb/master_pb/master.pb.go
+++ b/weed/pb/master_pb/master.pb.go
@@ -648,14 +648,15 @@ func (m *Location) GetPublicUrl() string {
 }
 
 type AssignRequest struct {
-	Count              uint64 `protobuf:"varint,1,opt,name=count" json:"count,omitempty"`
-	Replication        string `protobuf:"bytes,2,opt,name=replication" json:"replication,omitempty"`
-	Collection         string `protobuf:"bytes,3,opt,name=collection" json:"collection,omitempty"`
-	Ttl                string `protobuf:"bytes,4,opt,name=ttl" json:"ttl,omitempty"`
-	DataCenter         string `protobuf:"bytes,5,opt,name=data_center,json=dataCenter" json:"data_center,omitempty"`
-	Rack               string `protobuf:"bytes,6,opt,name=rack" json:"rack,omitempty"`
-	DataNode           string `protobuf:"bytes,7,opt,name=data_node,json=dataNode" json:"data_node,omitempty"`
-	MemoryMapMaxSizeMb uint32 `protobuf:"varint,8,opt,name=memory_map_max_size_mb,json=MemoryMapMaxSizeMb" json:"memory_map_max_size_mb,omitempty"`
+	Count               uint64 `protobuf:"varint,1,opt,name=count" json:"count,omitempty"`
+	Replication         string `protobuf:"bytes,2,opt,name=replication" json:"replication,omitempty"`
+	Collection          string `protobuf:"bytes,3,opt,name=collection" json:"collection,omitempty"`
+	Ttl                 string `protobuf:"bytes,4,opt,name=ttl" json:"ttl,omitempty"`
+	DataCenter          string `protobuf:"bytes,5,opt,name=data_center,json=dataCenter" json:"data_center,omitempty"`
+	Rack                string `protobuf:"bytes,6,opt,name=rack" json:"rack,omitempty"`
+	DataNode            string `protobuf:"bytes,7,opt,name=data_node,json=dataNode" json:"data_node,omitempty"`
+	MemoryMapMaxSizeMb  uint32 `protobuf:"varint,8,opt,name=memory_map_max_size_mb,json=MemoryMapMaxSizeMb" json:"memory_map_max_size_mb,omitempty"`
+	WritableVolumeCount uint32 `protobuf:"varint,9,opt,name=writable_volume_count" json:"writable_volume_count,omitempty"`
 }
 
 func (m *AssignRequest) Reset()                    { *m = AssignRequest{} }
diff --git a/weed/server/master_grpc_server_volume.go b/weed/server/master_grpc_server_volume.go
index 0580acf76..8fc56e9b8 100644
--- a/weed/server/master_grpc_server_volume.go
+++ b/weed/server/master_grpc_server_volume.go
@@ -78,7 +78,7 @@ func (ms *MasterServer) Assign(ctx context.Context, req *master_pb.AssignRequest
 		}
 		ms.vgLock.Lock()
 		if !ms.Topo.HasWritableVolume(option) {
-			if _, err = ms.vg.AutomaticGrowByType(option, ms.grpcDialOption, ms.Topo); err != nil {
+			if _, err = ms.vg.AutomaticGrowByType(option, ms.grpcDialOption, ms.Topo, int(req.WritableVolumeCount)); err != nil {
 				ms.vgLock.Unlock()
 				return nil, fmt.Errorf("Cannot grow volume group! %v", err)
 			}
diff --git a/weed/server/master_server_handlers.go b/weed/server/master_server_handlers.go
index 9bcd35ced..514d86800 100644
--- a/weed/server/master_server_handlers.go
+++ b/weed/server/master_server_handlers.go
@@ -100,6 +100,11 @@ func (ms *MasterServer) dirAssignHandler(w http.ResponseWriter, r *http.Request)
 		requestedCount = 1
 	}
 
+	writableVolumeCount, e := strconv.Atoi(r.FormValue("writableVolumeCount"))
+	if e != nil {
+		writableVolumeCount = 0
+	}
+
 	option, err := ms.getVolumeGrowOption(r)
 	if err != nil {
 		writeJsonQuiet(w, r, http.StatusNotAcceptable, operation.AssignResult{Error: err.Error()})
@@ -114,7 +119,7 @@ func (ms *MasterServer) dirAssignHandler(w http.ResponseWriter, r *http.Request)
 		ms.vgLock.Lock()
 		defer ms.vgLock.Unlock()
 		if !ms.Topo.HasWritableVolume(option) {
-			if _, err = ms.vg.AutomaticGrowByType(option, ms.grpcDialOption, ms.Topo); err != nil {
+			if _, err = ms.vg.AutomaticGrowByType(option, ms.grpcDialOption, ms.Topo, writableVolumeCount); err != nil {
 				writeJsonError(w, r, http.StatusInternalServerError,
 					fmt.Errorf("Cannot grow volume group! %v", err))
 				return
diff --git a/weed/topology/volume_growth.go b/weed/topology/volume_growth.go
index f21ab45ef..636eb2260 100644
--- a/weed/topology/volume_growth.go
+++ b/weed/topology/volume_growth.go
@@ -59,8 +59,11 @@ func (vg *VolumeGrowth) findVolumeCount(copyCount int) (count int) {
 	return
 }
 
-func (vg *VolumeGrowth) AutomaticGrowByType(option *VolumeGrowOption, grpcDialOption grpc.DialOption, topo *Topology) (count int, err error) {
-	count, err = vg.GrowByCountAndType(grpcDialOption, vg.findVolumeCount(option.ReplicaPlacement.GetCopyCount()), option, topo)
+func (vg *VolumeGrowth) AutomaticGrowByType(option *VolumeGrowOption, grpcDialOption grpc.DialOption, topo *Topology, targetCount int) (count int, err error) {
+	if targetCount == 0 {
+		targetCount = vg.findVolumeCount(option.ReplicaPlacement.GetCopyCount())
+	}
+	count, err = vg.GrowByCountAndType(grpcDialOption, targetCount, option, topo)
 	if count > 0 && count%option.ReplicaPlacement.GetCopyCount() == 0 {
 		return count, nil
 	}