From fad1ef7250e5b684d21e7c5c01a068292fcf59ad Mon Sep 17 00:00:00 2001 From: Nathan VanBenschoten Date: Thu, 30 Dec 2021 15:27:54 -0500 Subject: [PATCH] sql: use FastIntMap for QueryState.RangesPerNode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This commit switches `QueryState.RangesPerNode` from a `map[roachpb.NodeID]int` to a `util.FastIntMap`. This avoids a pair of heap allocation as long as node IDs say below 32 and the number of ranges for each node stays below 14. ``` name old time/op new time/op delta KV/Scan/SQL/rows=1-10 94.5µs ± 4% 94.1µs ± 6% ~ (p=0.796 n=10+10) name old alloc/op new alloc/op delta KV/Scan/SQL/rows=1-10 20.1kB ± 0% 19.9kB ± 0% -0.84% (p=0.000 n=10+9) name old allocs/op new allocs/op delta KV/Scan/SQL/rows=1-10 245 ± 0% 243 ± 0% -0.82% (p=0.000 n=10+9) ``` --- pkg/sql/physicalplan/replicaoracle/BUILD.bazel | 1 + pkg/sql/physicalplan/replicaoracle/oracle.go | 6 +++--- pkg/sql/physicalplan/span_resolver.go | 3 ++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/pkg/sql/physicalplan/replicaoracle/BUILD.bazel b/pkg/sql/physicalplan/replicaoracle/BUILD.bazel index 5a809a787762..4a38b3cc9a31 100644 --- a/pkg/sql/physicalplan/replicaoracle/BUILD.bazel +++ b/pkg/sql/physicalplan/replicaoracle/BUILD.bazel @@ -12,6 +12,7 @@ go_library( "//pkg/rpc", "//pkg/settings/cluster", "//pkg/sql/sqlerrors", + "//pkg/util", "@com_github_cockroachdb_errors//:errors", ], ) diff --git a/pkg/sql/physicalplan/replicaoracle/oracle.go b/pkg/sql/physicalplan/replicaoracle/oracle.go index 38eddb98d07e..67d375cda6f6 100644 --- a/pkg/sql/physicalplan/replicaoracle/oracle.go +++ b/pkg/sql/physicalplan/replicaoracle/oracle.go @@ -23,6 +23,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/rpc" "github.com/cockroachdb/cockroach/pkg/settings/cluster" "github.com/cockroachdb/cockroach/pkg/sql/sqlerrors" + "github.com/cockroachdb/cockroach/pkg/util" "github.com/cockroachdb/errors" ) @@ -106,14 +107,13 @@ var oracleFactories = map[Policy]OracleFactory{} // QueryState encapsulates the history of assignments of ranges to nodes // done by an oracle on behalf of one particular query. type QueryState struct { - RangesPerNode map[roachpb.NodeID]int + RangesPerNode util.FastIntMap AssignedRanges map[roachpb.RangeID]roachpb.ReplicaDescriptor } // MakeQueryState creates an initialized QueryState. func MakeQueryState() QueryState { return QueryState{ - RangesPerNode: make(map[roachpb.NodeID]int), AssignedRanges: make(map[roachpb.RangeID]roachpb.ReplicaDescriptor), } } @@ -233,7 +233,7 @@ func (o *binPackingOracle) ChoosePreferredReplica( minLoad := int(math.MaxInt32) var leastLoadedIdx int for i, repl := range replicas { - assignedRanges := queryState.RangesPerNode[repl.NodeID] + assignedRanges := queryState.RangesPerNode.GetDefault(int(repl.NodeID)) if assignedRanges != 0 && assignedRanges < o.maxPreferredRangesPerLeaseHolder { return repl.ReplicaDescriptor, nil } diff --git a/pkg/sql/physicalplan/span_resolver.go b/pkg/sql/physicalplan/span_resolver.go index ce47cda2f6e1..7138eeaaef88 100644 --- a/pkg/sql/physicalplan/span_resolver.go +++ b/pkg/sql/physicalplan/span_resolver.go @@ -267,7 +267,8 @@ func (it *spanResolverIterator) ReplicaInfo( if err != nil { return roachpb.ReplicaDescriptor{}, err } - it.queryState.RangesPerNode[repl.NodeID]++ + prev := it.queryState.RangesPerNode.GetDefault(int(repl.NodeID)) + it.queryState.RangesPerNode.Set(int(repl.NodeID), prev+1) it.queryState.AssignedRanges[rngID] = repl return repl, nil }