From afdc2e999e3c791c527ad66e7f1f665ef28928e8 Mon Sep 17 00:00:00 2001
From: Alex Sarkesian <sarkesian@cockroachlabs.com>
Date: Mon, 19 Dec 2022 18:00:11 -0500
Subject: [PATCH] server: implement decommission pre-check api

This change implements the `DecommissionPreCheck` RPC on the `Admin`
service, using the support for evaluating node decommission readiness by
checking each range introduced in #93758. In checking node decommission
readiness, only nodes that have a valid, non-`DECOMMISSIONED` liveness
status are checked, and ranges with replicas on the checked nodes that
encounter errors in attempting to allocate replacement replicas are
reported in the response. Ranges that have replicas on multiple checked
nodes have their errors reported for each nodeID in the request list.

Depends on #93758, #90222.

Epic: CRDB-20924

Release note: None
---
 pkg/server/admin.go | 83 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 82 insertions(+), 1 deletion(-)

diff --git a/pkg/server/admin.go b/pkg/server/admin.go
index 4a2a128d688a..b59b663989ff 100644
--- a/pkg/server/admin.go
+++ b/pkg/server/admin.go
@@ -2648,7 +2648,88 @@ func (s *adminServer) getStatementBundle(ctx context.Context, id int64, w http.R
 func (s *systemAdminServer) DecommissionPreCheck(
 	ctx context.Context, req *serverpb.DecommissionPreCheckRequest,
 ) (*serverpb.DecommissionPreCheckResponse, error) {
-	return nil, grpcstatus.Errorf(codes.Unimplemented, "method DecommissionPreCheck not implemented")
+	var collectTraces bool
+	if s := tracing.SpanFromContext(ctx); (s != nil && s.RecordingType() != tracingpb.RecordingOff) || req.CollectTraces {
+		collectTraces = true
+	}
+
+	// Initially evaluate node liveness status, so we filter the nodes to check.
+	var nodesToCheck []roachpb.NodeID
+	livenessStatusByNodeID, err := getLivenessStatusMap(ctx, s.nodeLiveness, s.clock.Now().GoTime(), s.st)
+	if err != nil {
+		return nil, serverError(ctx, err)
+	}
+
+	resp := &serverpb.DecommissionPreCheckResponse{}
+
+	// Any nodes that are already decommissioned or have unknown liveness should
+	// not be checked, and are added to response without replica counts or errors.
+	for _, nID := range req.NodeIDs {
+		livenessStatus := livenessStatusByNodeID[nID]
+		if livenessStatus == livenesspb.NodeLivenessStatus_UNKNOWN {
+			resp.CheckedNodes = append(resp.CheckedNodes, serverpb.DecommissionPreCheckResponse_NodeCheckResult{
+				NodeID:                nID,
+				DecommissionReadiness: serverpb.DecommissionPreCheckResponse_UNKNOWN,
+				LivenessStatus:        livenessStatus,
+			})
+		} else if livenessStatus == livenesspb.NodeLivenessStatus_DECOMMISSIONED {
+			resp.CheckedNodes = append(resp.CheckedNodes, serverpb.DecommissionPreCheckResponse_NodeCheckResult{
+				NodeID:                nID,
+				DecommissionReadiness: serverpb.DecommissionPreCheckResponse_ALREADY_DECOMMISSIONED,
+				LivenessStatus:        livenessStatus,
+			})
+		} else {
+			nodesToCheck = append(nodesToCheck, nID)
+		}
+	}
+
+	results, err := s.server.DecommissionPreCheck(ctx, nodesToCheck, req.StrictReadiness, collectTraces, int(req.NumReplicaReport))
+	if err != nil {
+		return nil, err
+	}
+
+	// Collect ranges that encountered errors by the nodes on which their replicas
+	// exist. Ranges with replicas on multiple checked nodes will result in the
+	// error being reported for each nodeID.
+	rangeCheckErrsByNode := make(map[roachpb.NodeID][]serverpb.DecommissionPreCheckResponse_RangeCheckResult)
+	for _, rangeWithErr := range results.rangesNotReady {
+		rangeCheckResult := serverpb.DecommissionPreCheckResponse_RangeCheckResult{
+			RangeID: rangeWithErr.desc.RangeID,
+			Action:  rangeWithErr.action,
+			Events:  recordedSpansToTraceEvents(rangeWithErr.tracingSpans),
+			Error:   rangeWithErr.err.Error(),
+		}
+
+		for _, nID := range nodesToCheck {
+			if rangeWithErr.desc.Replicas().HasReplicaOnNode(nID) {
+				rangeCheckErrsByNode[nID] = append(rangeCheckErrsByNode[nID], rangeCheckResult)
+			}
+		}
+	}
+
+	// Evaluate readiness for each node to check based on how many ranges have
+	// replicas on the node that did not pass checks.
+	for _, nID := range nodesToCheck {
+		numReplicas := len(results.replicasByNode[nID])
+		var readiness serverpb.DecommissionPreCheckResponse_NodeReadiness
+		if len(rangeCheckErrsByNode[nID]) > 0 {
+			readiness = serverpb.DecommissionPreCheckResponse_ALLOCATION_ERRORS
+		} else {
+			readiness = serverpb.DecommissionPreCheckResponse_READY
+		}
+
+		nodeCheckResult := serverpb.DecommissionPreCheckResponse_NodeCheckResult{
+			NodeID:                nID,
+			DecommissionReadiness: readiness,
+			LivenessStatus:        livenessStatusByNodeID[nID],
+			ReplicaCount:          int64(numReplicas),
+			CheckedRanges:         rangeCheckErrsByNode[nID],
+		}
+
+		resp.CheckedNodes = append(resp.CheckedNodes, nodeCheckResult)
+	}
+
+	return resp, nil
 }
 
 // DecommissionStatus returns the DecommissionStatus for all or the given nodes.