complex sticky: remove A* and just use breadth first

After further thought, the heuristic function was not entirely accurate
yet. Our search was operating on # of steals, but the heuristic was
returning how beneficial it would be to steal a level. A more accurate
guess from any node would be 1 unless the node happens to be a target
node. I benchmarked that and it was marginally better than the previous
code, but just removing the heuristic leads to a better improvement and
keeps things simpler, so we may as well just remove the heuristic.

Removing the heuristic allows us to remove fscore and simplifies the
heap comparison logic, which actually gives us a speed boost.

name                           old time/op    new time/op    delta
Large-8                          2.15ms ± 1%    2.15ms ± 1%    ~     (p=0.968 n=9+10)
LargeWithExisting-8              8.67ms ± 1%    8.64ms ± 2%    ~     (p=0.079 n=9+10)
LargeImbalanced-8                2.93ms ± 4%    2.88ms ± 2%    ~     (p=0.089 n=10+10)
LargeWithExistingImbalanced-8    8.75ms ± 1%    8.69ms ± 1%  -0.65%  (p=0.002 n=9+9)
Java/large-8                      159ms ± 1%     158ms ± 1%  -0.76%  (p=0.023 n=10+10)
Java/large_imbalance-8            197ms ± 3%     188ms ± 2%  -4.38%  (p=0.000 n=10+9)
Java/medium-8                    8.94ms ± 1%    8.86ms ± 1%    ~     (p=0.053 n=10+9)
Java/medium_imbalance-8          10.5ms ± 2%    10.3ms ± 1%  -2.38%  (p=0.000 n=10+8)
Java/small-8                     6.81ms ± 2%    6.80ms ± 1%    ~     (p=0.604 n=10+9)
Java/small_imbalance-8           8.43ms ± 1%    8.27ms ± 1%  -1.94%  (p=0.000 n=9+6)

name                           old alloc/op   new alloc/op   delta
Large-8                          1.92MB ± 0%    1.92MB ± 0%    ~     (p=0.060 n=10+10)
LargeWithExisting-8              3.69MB ± 0%    3.69MB ± 0%    ~     (p=0.661 n=9+10)
LargeImbalanced-8                2.12MB ± 1%    2.12MB ± 1%    ~     (p=0.912 n=10+10)
LargeWithExistingImbalanced-8    3.72MB ± 0%    3.72MB ± 0%    ~     (p=0.071 n=8+10)
Java/large-8                      127MB ± 0%     127MB ± 0%    ~     (p=0.069 n=10+9)
Java/large_imbalance-8            131MB ± 0%     131MB ± 0%  -0.01%  (p=0.000 n=8+9)
Java/medium-8                    7.78MB ± 0%    7.78MB ± 0%    ~     (p=0.210 n=10+10)
Java/medium_imbalance-8          8.06MB ± 0%    8.05MB ± 0%  -0.10%  (p=0.000 n=10+10)
Java/small-8                     6.20MB ± 0%    6.20MB ± 0%    ~     (p=0.109 n=10+10)
Java/small_imbalance-8           6.42MB ± 0%    6.42MB ± 0%  -0.10%  (p=0.000 n=10+10)

name                           old allocs/op  new allocs/op  delta
Large-8                             335 ± 0%       335 ± 0%    ~     (p=0.961 n=10+8)
LargeWithExisting-8               18.5k ± 1%     18.6k ± 1%    ~     (p=0.119 n=10+10)
LargeImbalanced-8                   785 ± 7%       792 ± 7%    ~     (p=0.631 n=10+10)
LargeWithExistingImbalanced-8     18.4k ± 1%     18.5k ± 0%    ~     (p=0.085 n=10+10)
Java/large-8                      6.04k ± 0%     6.04k ± 0%    ~     (all equal)
Java/large_imbalance-8            7.43k ± 0%     7.43k ± 0%    ~     (p=0.678 n=8+9)
Java/medium-8                     3.03k ± 0%     3.03k ± 0%    ~     (all equal)
Java/medium_imbalance-8           3.15k ± 0%     3.15k ± 0%    ~     (p=0.670 n=10+10)
Java/small-8                      2.46k ± 0%     2.46k ± 0%    ~     (p=0.108 n=10+9)
Java/small_imbalance-8            2.55k ± 0%     2.55k ± 0%    ~     (p=0.493 n=10+10)

Author: twmb

pkg/kgo/internal/sticky/graph.go

@@ -16,7 +16,7 @@ type graph struct {
 	// edge => who owns this edge; built in balancer's assignUnassigned
 	cxns []uint16
 
-	// scores are all node scores from a seach node. The gscore field
+	// scores are all node scores from a seach node. The distance field
 	// is reset on findSteal to noScore.
 	scores pathScores
 
@@ -59,19 +59,18 @@ func (g *graph) changeOwnership(edge int32, newDst uint16) {
 	g.cxns[edge] = newDst
 }
 
-// findSteal uses A* search to find a path from the best node it can reach.
+// findSteal uses Dijkstra search to find a path from the best node it can reach.
 func (g *graph) findSteal(from uint16) ([]stealSegment, bool) {
 	// First, we must reset our scores from any prior run. This is O(M),
 	// but is fast and faster than making a map and extending it a lot.
 	for i := range g.scores {
-		g.scores[i].gscore = noScore
+		g.scores[i].distance = noScore
 		g.scores[i].done = false
 	}
 
 	first, _ := g.getScore(from)
 
-	first.gscore = 0
-	first.fscore = h(first, first)
+	first.distance = 0
 	first.done = true
 
 	g.heapBuf = append(g.heapBuf[:0], first)
@@ -104,22 +103,22 @@ func (g *graph) findSteal(from uint16) ([]stealSegment, bool) {
 					continue
 				}
 
-				gscore := current.gscore + 1
-				// If our neghbor gscore is less or equal, then we can
+				distance := current.distance + 1
+				// If our neghbor distance is less or equal, then we can
 				// reach the neighbor through a previous route we have
 				// tried and should not try again.
-				if gscore < neighbor.gscore {
+				if distance < neighbor.distance {
 					neighbor.parent = current
 					neighbor.srcEdge = edge
-					neighbor.gscore = gscore
-					neighbor.fscore = gscore + h(first, neighbor)
+					neighbor.distance = distance
 					if isNew {
 						heap.Push(rem, neighbor)
 					}
+
 					// We never need to fix the heap position.
-					// Our level and fscore is static, and once
-					// we set gscore, it is the minumum it will be
-					// and we never revisit this neighbor.
+					// Our level is static, and once we set
+					// distance, it is the minimum it will be
+					// and we never revisit the neighbor.
 				}
 			}
 		}
@@ -135,43 +134,27 @@ type stealSegment struct {
 }
 
 type pathScore struct {
-	done    bool
-	node    uint16 // member num
-	parent  *pathScore
-	srcEdge int32 // partNum
-	level   int32 // partitions owned on this segment
-	gscore  int32 // how many steals it would take to get here
-	fscore  int32
+	done     bool
+	node     uint16 // member num
+	distance int32  // how many steals it would take to get here
+	srcEdge  int32  // partNum
+	level    int32  // partitions owned on this segment
+	parent   *pathScore
 }
 
 type pathScores []pathScore
 
 const infinityScore = 1<<31 - 1
 const noScore = -1
 
-// For A*, if we never overestimate (with h), then the path we find is
-// optimal. A true estimation of our distance to any node is the node's
-// level minus ours.
-//
-// At worst, our target must be +2 levels from us. So, our estimation
-// any node can be our level, +2, minus theirs.
-func h(first, target *pathScore) int32 {
-	r := first.level + 2 - target.level
-	if r < 0 {
-		return 0
-	}
-	return r
-}
-
 func (g *graph) getScore(node uint16) (*pathScore, bool) {
 	r := &g.scores[node]
-	exists := r.gscore != noScore
+	exists := r.distance != noScore
 	if !exists {
 		*r = pathScore{
-			node:   node,
-			level:  int32(len(g.b.plan[node])),
-			gscore: infinityScore,
-			fscore: infinityScore,
+			node:     node,
+			level:    int32(len(g.b.plan[node])),
+			distance: infinityScore,
 		}
 	}
 	return r, !exists
@@ -188,9 +171,8 @@ func (p *pathHeap) Swap(i, j int) {
 func (p *pathHeap) Less(i, j int) bool {
 	l, r := (*p)[i], (*p)[j]
 	return l.level > r.level || l.level == r.level &&
-		(l.fscore < r.fscore || l.fscore == r.fscore &&
-			(l.gscore < r.gscore || l.gscore == r.gscore &&
-				l.node < r.node))
+		(l.distance < r.distance || l.distance == r.distance &&
+			l.node < r.node)
 }
 
 func (p *pathHeap) Push(x interface{}) { *p = append(*p, x.(*pathScore)) }