1import { assert, describe, it } from "@effect/vitest"
2import * as Hash from "effect/Hash"
3import * as HashRing from "effect/HashRing"
4import * as PrimaryKey from "effect/PrimaryKey"
5
6describe("HashRing", () => {
7 it("updates the stored node when adding the same primary key", () => {
8 const first = {
9 name: "first",
10 [PrimaryKey.symbol]() {
11 return "node"
12 }
13 }
14 const updated = {
15 name: "updated",
16 [PrimaryKey.symbol]() {
17 return "node"
18 }
19 }
20 const ring = HashRing.make<typeof first>()
21
22 HashRing.add(ring, first)
23 HashRing.add(ring, updated)
24
25 assert.strictEqual(HashRing.get(ring, "request"), updated)
26 })
27
28 it("updates the stored node when changing its weight", () => {
29 const first = {
30 name: "first",
31 [PrimaryKey.symbol]() {
32 return "node"
33 }
34 }
35 const updated = {
36 name: "updated",
37 [PrimaryKey.symbol]() {
38 return "node"
39 }
40 }
41 const ring = HashRing.make<typeof first>()
42
43 HashRing.add(ring, first)
44 HashRing.add(ring, updated, { weight: 2 })
45
46 assert.strictEqual(HashRing.get(ring, "request"), updated)
47 })
48
49 it("getShards normalizes the shard count", () => {
50 const node = {
51 name: "node",
52 [PrimaryKey.symbol]() {
53 return "node"
54 }
55 }
56 const ring = HashRing.make<typeof node>()
57 HashRing.add(ring, node)
58
59 const shards = [Number.NaN, -1, 0, 0.5, 1.9, 2.9].map((count) => HashRing.getShards(ring, count))
60
61 assert.deepStrictEqual(shards, [[], [], [], [], [node], [node, node]])
62 })
63
64 it("getShards considers the first ring entry when excluding allocated nodes", () => {
65 // With one ring entry per node, these keys make every shard nearest to
66 // `second`, so `first` at index 0 is reached only by the exclusion scan.
67 const first = {
68 [PrimaryKey.symbol]() {
69 return "node-10"
70 }
71 }
72 const second = {
73 [PrimaryKey.symbol]() {
74 return "node-29"
75 }
76 }
77 const ring = HashRing.make<typeof first>({ baseWeight: 1 })
78 HashRing.add(ring, first)
79 HashRing.add(ring, second)
80
81 assert.deepStrictEqual(HashRing.getShards(ring, 3)?.map(PrimaryKey.value), ["node-10", "node-29", "node-29"])
82 })
83
84 it("getShards is unchanged after adding and removing a fractional-weight node", () => {
85 const [a, b, c, d] = ["runner-a:34431", "runner-b:34431", "runner-c:34431", "runner-d:34431"].map(makeNode)
86 // Adding and removing `d` leaves a different floating-point weight total
87 // than adding only `a`, `b` and `c`.
88 const churned = HashRing.make<Node>()
89 HashRing.addMany(churned, [a, b, c], { weight: 0.1 })
90 HashRing.add(churned, d, { weight: 0.3 })
91 HashRing.remove(churned, d)
92 const fresh = HashRing.make<Node>()
93 HashRing.addMany(fresh, [a, b, c], { weight: 0.1 })
94
95 assert.deepStrictEqual(HashRing.getShards(churned, 300), HashRing.getShards(fresh, 300))
96 })
97
98 it("getShards breaks ring hash ties independently of insertion order", () => {
99 const first = makeNode("runner-27:34431")
100 const second = makeNode("runner-1176:34431")
101 // These virtual points of different nodes have the same hash.
102 assert.strictEqual(Hash.string("runner-27:34431:13"), Hash.string("runner-1176:34431:71"))
103 const ring = HashRing.make<Node>()
104 HashRing.add(ring, first)
105 HashRing.add(ring, second)
106 const swapped = HashRing.make<Node>()
107 HashRing.add(swapped, second)
108 HashRing.add(swapped, first)
109
110 assert.deepStrictEqual(
111 HashRing.getShards(swapped, 300),
112 HashRing.getShards(ring, 300)
113 )
114 })
115
116 it("getShards keeps the assignment for integer weights", () => {
117 const ring = HashRing.make<Node>()
118 const nodes = Array.from({ length: 5 }, (_, i) => makeNode(`runner-${i}:34431`))
119 const weights = [1, 1, 2, 1, 3]
120 weights.forEach((weight, i) => HashRing.add(ring, nodes[i], { weight }))
121
122 // Assignment from the implementation before deterministic ordering.
123 assert.deepStrictEqual(
124 HashRing.getShards(ring, 16),
125 [4, 3, 0, 2, 2, 4, 3, 4, 1, 2, 2, 4, 1, 4, 4, 0].map((i) => nodes[i])
126 )
127 })
128})
129
130interface Node extends PrimaryKey.PrimaryKey {}
131
132const makeNode = (key: string): Node => ({
133 [PrimaryKey.symbol]() {
134 return key
135 }
136})