From 04f4b108842742a7b0205a02baed32d3bacbb73f Mon Sep 17 00:00:00 2001
From: Konstantin Lebedev <9497591+kmlebedev@users.noreply.github.com>
Date: Fri, 12 Jul 2024 23:19:08 +0500
Subject: [PATCH] fix: avoid timeout if datacenter does not exist in topology
 (#5772)

* fix: avoid timeout if datacenter does not exist in topology

* fix: error msg

* fix: rm dublicate check

* fix: compare

* revert minor change
---
 weed/server/master_grpc_server_assign.go    |  4 ++++
 weed/server/master_grpc_server_volume.go    |  2 ++
 weed/server/master_server_handlers.go       |  7 +++++++
 weed/server/master_server_handlers_admin.go |  2 ++
 weed/topology/topology.go                   | 16 ++++++++++++++++
 5 files changed, 31 insertions(+)

diff --git a/weed/server/master_grpc_server_assign.go b/weed/server/master_grpc_server_assign.go
index 7455c9ea4..523db6038 100644
--- a/weed/server/master_grpc_server_assign.go
+++ b/weed/server/master_grpc_server_assign.go
@@ -69,6 +69,10 @@ func (ms *MasterServer) Assign(ctx context.Context, req *master_pb.AssignRequest
 		MemoryMapMaxSizeMb: req.MemoryMapMaxSizeMb,
 	}
 
+	if !ms.Topo.DataCenterExists(option.DataCenter) {
+		return nil, fmt.Errorf("data center %v not found in topology", option.DataCenter)
+	}
+
 	vl := ms.Topo.GetVolumeLayout(option.Collection, option.ReplicaPlacement, option.Ttl, option.DiskType)
 
 	var (
diff --git a/weed/server/master_grpc_server_volume.go b/weed/server/master_grpc_server_volume.go
index 86b66a818..afe062a2d 100644
--- a/weed/server/master_grpc_server_volume.go
+++ b/weed/server/master_grpc_server_volume.go
@@ -59,6 +59,8 @@ func (ms *MasterServer) ProcessGrowRequest() {
 						for _, newVidLocation := range newVidLocations {
 							ms.broadcastToClients(&master_pb.KeepConnectedResponse{VolumeLocation: newVidLocation})
 						}
+					} else {
+						glog.V(1).Infof("automatic volume grow failed: %+v", err)
 					}
 					vl.DoneGrowRequest()
 
diff --git a/weed/server/master_server_handlers.go b/weed/server/master_server_handlers.go
index e4188420d..65c4589d8 100644
--- a/weed/server/master_server_handlers.go
+++ b/weed/server/master_server_handlers.go
@@ -126,6 +126,13 @@ func (ms *MasterServer) dirAssignHandler(w http.ResponseWriter, r *http.Request)
 		startTime  = time.Now()
 	)
 
+	if !ms.Topo.DataCenterExists(option.DataCenter) {
+		writeJsonQuiet(w, r, http.StatusBadRequest, operation.AssignResult{
+			Error: fmt.Sprintf("data center %v not found in topology", option.DataCenter),
+		})
+		return
+	}
+
 	for time.Now().Sub(startTime) < maxTimeout {
 		fid, count, dnList, shouldGrow, err := ms.Topo.PickForWrite(requestedCount, option, vl)
 		if shouldGrow && !vl.HasGrowRequest() {
diff --git a/weed/server/master_server_handlers_admin.go b/weed/server/master_server_handlers_admin.go
index f40b819af..07c0a94d2 100644
--- a/weed/server/master_server_handlers_admin.go
+++ b/weed/server/master_server_handlers_admin.go
@@ -81,6 +81,8 @@ func (ms *MasterServer) volumeGrowHandler(w http.ResponseWriter, r *http.Request
 	if count, err = strconv.Atoi(r.FormValue("count")); err == nil {
 		if ms.Topo.AvailableSpaceFor(option) < int64(count*option.ReplicaPlacement.GetCopyCount()) {
 			err = fmt.Errorf("only %d volumes left, not enough for %d", ms.Topo.AvailableSpaceFor(option), count*option.ReplicaPlacement.GetCopyCount())
+		} else if !ms.Topo.DataCenterExists(option.DataCenter) {
+			err = fmt.Errorf("data center %v not found in topology", option.DataCenter)
 		} else {
 			var newVidLocations []*master_pb.VolumeLocation
 			newVidLocations, err = ms.vg.GrowByCountAndType(ms.grpcDialOption, count, option, ms.Topo)
diff --git a/weed/topology/topology.go b/weed/topology/topology.go
index 0a4cb4050..01822cbf2 100644
--- a/weed/topology/topology.go
+++ b/weed/topology/topology.go
@@ -285,6 +285,22 @@ func (t *Topology) UnRegisterVolumeLayout(v storage.VolumeInfo, dn *DataNode) {
 	}
 }
 
+func (t *Topology) DataCenterExists(dcName string) bool {
+	return dcName == "" || t.GetOrCreateDataCenter(dcName) != nil
+}
+
+func (t *Topology) GetDataCenter(dcName string) (dc *DataCenter) {
+	t.RLock()
+	defer t.RUnlock()
+	for _, c := range t.children {
+		dc = c.(*DataCenter)
+		if string(dc.Id()) == dcName {
+			return dc
+		}
+	}
+	return dc
+}
+
 func (t *Topology) GetOrCreateDataCenter(dcName string) *DataCenter {
 	t.Lock()
 	defer t.Unlock()