Merge #152986 #153059 #153061 #153160 #153281

152986: workload: set allow_unsafe_internals for workload commands r=alyshanjahani-crl a=alyshanjahani-crl

Previously some workload `run` and `init` commands would fail if and when they try to access crdb_internal or system tables (if the `allow_unsafe_internals` session var is globally defaulted to false - which will be the case in 26.1)

This commit modifies all connections that workload commands use to have the `allow_unsafe_internals` session variable set to on.

Fixes: #152755
Release note: None

153059: jsonpath: add err check for jsonValue Iterator retrieval r=ZhouXing19 a=ZhouXing19

There are 2 types that can return `ObjectJSONType` as the type: `json.jsonObject` or `json.jsonEncoded`. The latter can return not-nil error in their implementation of `ObjectIter()`. This commit is to add error check for this case.

Epic: none
Release note (bug fix): adding error check for jsonValue Iterator when evaluating a jsonpath on a json value.

153061: jsonpath/eval: graceful type unmatched handling r=ZhouXing19 a=ZhouXing19

Previously, if the type is not found in the switch case, we panic with a unhandled type error. Now we implement a more graceful way for erroring out.

Epic: none
Release note: None

153160: allocator: add doc.md r=tbg a=wenyihu6

This commit adds a notes page covering some parts of the allocator code from
my code reading.

Epic: none 
Release note: none 

153281: cdctest/validator: handle error early r=wenyihu6 a=wenyihu6

This commit properly handles errors from `rowDatums.FetchValKey`.

Epic: none
Release note: none

Co-authored-by: Alyshan Jahani <[email protected]>
Co-authored-by: ZhouXing19 <[email protected]>
Co-authored-by: wenyihu6 <[email protected]>

Author: 4 people

pkg/ccl/changefeedccl/cdctest/validator.go

@@ -321,14 +321,13 @@ func (v *beforeAfterValidator) checkRowAt(
 				stmtBuf.WriteString(` AND `)
 			}
 			jsonCol, err := rowDatums.FetchValKey(col)
-
+			if err != nil {
+				return err
+			}
 			if jsonCol == nil || jsonCol.Type() == json.NullJSONType {
 				fmt.Fprintf(&stmtBuf, `%s IS NULL`, col)
 			} else {
 				fmt.Fprintf(&stmtBuf, `to_json(%s)::TEXT = $%d`, col, i+1)
-				if err != nil {
-					return err
-				}
 				args = append(args, jsonCol.String())
 			}
 		}

pkg/kv/kvserver/allocator/doc.md

@@ -0,0 +1,46 @@
+# Allocator
+
+### Replica Scanner And Base Queue
+
+- Every store has a replica scanner that periodically scans and iterates over every replica of stores. It calls every base queues’s [`MaybeAddAsync`](https://github.com/cockroachdb/cockroach/blob/308665f4fb5176f8f7dcefdbfe8b9eee669e9f5a/pkg/kv/kvserver/queue.go#L635) with every replica. Every base queue’s `MaybeAdd` then calls into base queue’s `replicaCanBeProcessed`, shouldQueue. If `shouldQueue` returns true, it [adds](https://github.com/cockroachdb/cockroach/blob/308665f4fb5176f8f7dcefdbfe8b9eee669e9f5a/pkg/kv/kvserver/queue.go#L763) the replica to the queue with the priority given by `shouldQueue`. Skips if `shouldQueue` returns false.
+- Every base queue implements `bq.processLoop` which pops items added by the replica scanner and calls into `bq.processOneAsyncAndReleaseSem`. `bq.processOneAsyncAndReleaseSem` then calls `bq.processReplica` async with a time out. `bq.processReplica` then calls `bq.process`. If replicas [fail](https://github.com/cockroachdb/cockroach/blob/308665f4fb5176f8f7dcefdbfe8b9eee669e9f5a/pkg/kv/kvserver/queue.go#L1141-L1143) to be processed, they will be sent to the purgatory queue.
+- Base queue includes: consistency queue, lease queue, merge queue, mvcc gc queue, raft log queue, raft snapshot queue, replica gc queue, replicate queue, split queue, timeseries maintenance queue.
+
+
+## Replicate Queue
+### High Level
+1. `ReplicateQueue.shouldQueue` calls into `ReplicaPlanner.ShouldPlanChange` which decides whether the replica scanner adds the replica to the queue.
+2. Replica scanner adds replicas that need to be queued to the base queue.
+3. Inside `replicateQueue.process` with one replica, a) grab the range token b) calls into `rq.processOneChange`.
+4. `processOneChange` calls into `ReplicaPlanner.PlanOneChange` which plans a replicate change or an error. `PlanOneChange` calls into `allocator.ComputeAction` which returns an enum action. This enum action tells replicate queue to consider add / remove / replace replicas. `ComputeAction` gives an action that should be done.
+   The separation is here because `planner.ShouldPlanChange` calls into `ComputeAction` as well, and we want a cheap check for the queue. Note that actions returned from `ComputeAction` in `ShouldPlanChange` v.s. `PlanOneChange` may be different.
+5. Replica planner then switches and computes what replicate changes need to be done based on the action.
+
+### shouldQueue: whether queue replicas to replicate queue
+1. `shouldQueue` takes a replica and calls into [`ReplicaPlanner.ShouldPlanChange`](https://github.com/cockroachdb/cockroach/blob/5c2a22f84f0581df39239bc890b9d7d08c6e020a/pkg/kv/kvserver/allocator/plan/replicate.go#L145) which then calls into the allocator via `rp.allocator.ComputeAction`.
+2. [`ComputeAction`](https://github.com/cockroachdb/cockroach/blob/ad4d1478461f2edfa49d1cadd34cad910f5f5b10/pkg/kv/kvserver/allocator/allocatorimpl/allocator.go#L913) computes an action that needs to be done on this replica.
++ It first checks if a repair action is needed: under-replicated, quorum, dead/decommissioning replicas, over-replicated. Note that this does not include constraint repair.
++ If no repair needs to be done on this range, it would fall back and return `AllocatorConsiderRebalance` by default.
+3. After ComputeAction gives `ReplicaPlanner.ShouldPlanChange` the action, it would return true if a repair action is needed. Otherwise, it checks if a rebalance target can be found via [Allocator.RebalanceTarget](https://github.com/cockroachdb/cockroach/blob/ad4d1478461f2edfa49d1cadd34cad910f5f5b10/pkg/kv/kvserver/allocator/allocatorimpl/allocator.go#L1784). We will go into the details of `RebalanceTarget` below.
+
+### processOneChange: plan and process a change on the replica
+1. [`replicateQueue.processOneChange`](https://github.com/cockroachdb/cockroach/blob/a96cb747d2043d2959c1d44e29064c57d7aacde8/pkg/kv/kvserver/replicate_queue.go#L941) happens when the base queue pops the replica off the queue and actually start processing the replica.
+2. `replicateQueue.processOneChange` calls into `ReplicaPlanner.PlanOneChange` which then calls into `ComputeAction` as well. Based on the action, it would call into ReplicaPlanner's helper functions (such as [`addOrReplaceVoters`](https://github.com/cockroachdb/cockroach/blob/5c2a22f84f0581df39239bc890b9d7d08c6e020a/pkg/kv/kvserver/allocator/plan/replicate.go#L355)) to compute the exact allocation operation out (such as a lease transfer, change of replicas).
+3. For now, lets focus on how [`considerRebalance`](https://github.com/cockroachdb/cockroach/blob/5c2a22f84f0581df39239bc890b9d7d08c6e020a/pkg/kv/kvserver/allocator/plan/replicate.go#L773). As a reminder, this action is returned by default for every replica when no other repair action applies. Note that this action does help repair zone constraints.
+
+### considerRebalance: find one rebalancing target for any of the existing replica
+1. `ReplicaPlanner` first [picks](https://github.com/cockroachdb/cockroach/blob/5c2a22f84f0581df39239bc890b9d7d08c6e020a/pkg/kv/kvserver/allocator/plan/replicate.go#L789) which scorer options to be used for the allocator decision. It uses the `RangeCountScorerOptions` by default and `ScatterScorerOptions` if this action originates from a scatter operation. [ScorerOptions](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L266) define the heuristics used to evaluate candidate replicas, helping the allocator guide decisions toward specific goals such as range count convergence.
+2. `ReplicaPlanner` then calls into the allocator to find the best rebalancing target via [`Allocator.RebalanceTarget`](https://github.com/cockroachdb/cockroach/blob/ad4d1478461f2edfa49d1cadd34cad910f5f5b10/pkg/kv/kvserver/allocator/allocatorimpl/allocator.go#L1784).
+3. `Allocator.RebalanceTarget` first [builds](https://github.com/cockroachdb/cockroach/blob/ad4d1478461f2edfa49d1cadd34cad910f5f5b10/pkg/kv/kvserver/allocator/allocatorimpl/allocator.go#L1818-L1820) constraint checkers, which are passed to rankedCandidateListForRebalancing to determine which existing replicas / stores are valid or required based on the range’s constraints.
+4. `Allocator.RebalanceTarget` then calls into [`rankedCandidateListForRebalancing`](https://github.com/cockroachdb/cockroach/blob/ad4d1478461f2edfa49d1cadd34cad910f5f5b10/pkg/kv/kvserver/allocator/allocatorimpl/allocator.go#L1848) to find a list of rebalance options. Each rebalance option consists of an existing replica and a set of candidate stores that can replace it.
+- `rankedCandidateListForRebalancing` first [builds](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1513) a map of existing store ids to their corresponding candidate entries - annotated with attributes such as valid, necessary, disk fullness, I/O overload, and diversity score.
+- For each existing replica store, it then constructs an [`equivalence class`](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1579), containing the store itself along with a list of comparable candidates.
+  - This is how each equivalence class is constraucted: for every existing replica, all other stores are considered as potential candidates. During this iteration, some [filtering](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1650) is applied based on the state of other existing replicas. If a candidate store is [not worse](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1671) than the existing store, it is added to the list of comparable candidates. From all considered candidates, the best set is selected, and an equivalance class is constructed using existing replica and these best candidates. Note that at this stage, only attributes such as valid, necessary (constraint conformance), diversity score, and disk fullness are considered. At this point, range count convergence has not been considered yet.
+
+  - (TODO(wenyihu6 during review): why ioOverloaded is always false here there is a comment https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1550-L1553 but idu).
+
+  - (TODO(wenyihu6 during review) looks like at this stage we do not exclude other existing replicas stores - we filter it out below when iterating over comparable candidates but why)
+
+- At this point, we have an equivalence class for every existing replica. We then examine candidates’ range count statistics within each class to help break ties.
+  - For each equivalence class, ScorerOptions [populate](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1775) the convergence score, balance score, and range count for existing store and potential candidates. The candidate set is then further filtered to include only those [better](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1792) than the existing store.
+5. Using the results from rankedCandidateListForRebalancing, `Allocator.RebalanceTarget` iterates over all rebalanceOptions to [identify](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1826) the option that offers the largest improvement. As a reminder, each `rebalanceOption` consists of an existing replica and a list of candidate stores. For each option, it selects the best candidate set, randomly [chooses](https://github.com/cockroachdb/cockroach/blob/bc9fdcd029eae1f30a5ea82e44b60a629d452f6c/pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go#L1821) one candidate, and then determines the overall best rebalance option by comparing all rebalanceOptions.

pkg/util/jsonpath/eval/eval.go

@@ -267,14 +267,17 @@ func (ctx *jsonpathCtx) executeAnyItem(
 			}
 		}
 	case json.ObjectJSONType:
-		iter, _ := jsonValue.ObjectIter()
+		iter, err := jsonValue.ObjectIter()
+		if err != nil {
+			return nil, errors.Wrapf(err, "getting iterator for json object")
+		}
 		for iter.Next() {
 			if err := processItem(iter.Value()); err != nil {
 				return nil, err
 			}
 		}
 	default:
-		panic(errors.AssertionFailedf("executeAnyItem called with type: %s", jsonValue.Type()))
+		return nil, errors.AssertionFailedf("executeAnyItem called with type: %s", jsonValue.Type())
 	}
 	return agg, nil
 }

pkg/util/jsonpath/eval/method.go

@@ -83,7 +83,7 @@ func (ctx *jsonpathCtx) evalNumericMethod(
 			dec.Negative = false
 		}
 	default:
-		panic(errors.Newf("unimplemented method: %s", method.Type))
+		return nil, errors.Newf("unsupported jsonpath method type %d for numeric evaluation", method.Type)
 	}
 	if err != nil {
 		return nil, err

pkg/util/jsonpath/eval/operation.go

@@ -75,7 +75,7 @@ func (ctx *jsonpathCtx) evalOperation(
 	case jsonpath.OpPlus, jsonpath.OpMinus:
 		return ctx.evalUnaryArithmetic(op, jsonValue)
 	default:
-		panic(errors.AssertionFailedf("unhandled operation type"))
+		return nil, errors.AssertionFailedf("unhandled operation type")
 	}
 }
 
@@ -98,7 +98,7 @@ func (ctx *jsonpathCtx) evalBoolean(
 	case jsonpath.OpStartsWith:
 		return ctx.evalPredicate(op, jsonValue, evalStartsWithFunc, true /* evalRight */, false /* unwrapRight */)
 	default:
-		panic(errors.AssertionFailedf("unhandled operation type"))
+		return jsonpathBoolUnknown, errors.AssertionFailedf("unhandled operation type")
 	}
 }
 
@@ -224,7 +224,7 @@ func (ctx *jsonpathCtx) evalLogical(
 		}
 		return jsonpathBoolTrue, nil
 	default:
-		panic(errors.AssertionFailedf("unhandled logical operation type"))
+		return jsonpathBoolUnknown, errors.AssertionFailedf("unhandled logical operation type")
 	}
 
 	rightOp, ok := op.Right.(jsonpath.Operation)
@@ -247,7 +247,7 @@ func (ctx *jsonpathCtx) evalLogical(
 		}
 		return rightBool, nil
 	default:
-		panic(errors.AssertionFailedf("unhandled logical operation type"))
+		return jsonpathBoolUnknown, errors.AssertionFailedf("unhandled logical operation type")
 	}
 }
 
@@ -359,7 +359,7 @@ func evalComparisonFunc(operation jsonpath.Operation, l, r json.JSON) (jsonpathB
 		// Don't evaluate non-scalar types.
 		return jsonpathBoolUnknown, nil
 	default:
-		panic(errors.AssertionFailedf("unhandled json type"))
+		return jsonpathBoolUnknown, errors.AssertionFailedf("unhandled json type")
 	}
 
 	var res bool
@@ -377,7 +377,7 @@ func evalComparisonFunc(operation jsonpath.Operation, l, r json.JSON) (jsonpathB
 	case jsonpath.OpCompGreaterEqual:
 		res = cmp >= 0
 	default:
-		panic(errors.AssertionFailedf("unhandled jsonpath comparison type"))
+		return jsonpathBoolUnknown, errors.AssertionFailedf("unhandled jsonpath comparison type")
 	}
 	if res {
 		return jsonpathBoolTrue, nil
@@ -473,7 +473,7 @@ func performArithmetic(
 		}
 		_, err = tree.DecimalCtx.Rem(&res, leftNum, rightNum)
 	default:
-		panic(errors.AssertionFailedf("unhandled jsonpath arithmetic type"))
+		return nil, errors.AssertionFailedf("unhandled jsonpath arithmetic type")
 	}
 	if err != nil {
 		return nil, err

pkg/workload/connection.go

@@ -118,6 +118,7 @@ func SanitizeUrls(gen Generator, connFlags *ConnFlags, urls []string) (string, e
 func SetUrlConnVars(gen Generator, connFlags *ConnFlags, urls []string) error {
 	vars := make(map[string]string)
 	vars["application_name"] = gen.Meta().Name
+	vars["allow_unsafe_internals"] = "true"
 	if connFlags != nil {
 		if connFlags.IsoLevel != "" {
 			// As a convenience, replace underscores with spaces. This allows users of