|
| 1 | +package worker |
| 2 | + |
| 3 | +import ( |
| 4 | + "time" |
| 5 | + |
| 6 | + . "github.com/onsi/ginkgo/v2" |
| 7 | + . "github.com/onsi/gomega" |
| 8 | +) |
| 9 | + |
| 10 | +// mustAllocate fails the spec if allocation errors, so the ordinary "give me a |
| 11 | +// port" specs below stay readable while still refusing to assert on a zero |
| 12 | +// value that the allocator never intended to hand out. |
| 13 | +func mustAllocate(s *backendSupervisor, key string) int { |
| 14 | + port, err := s.allocatePort(key) |
| 15 | + ExpectWithOffset(1, err).NotTo(HaveOccurred()) |
| 16 | + return port |
| 17 | +} |
| 18 | + |
| 19 | +var _ = Describe("Backend gRPC port allocator", func() { |
| 20 | + Describe("range exhaustion", func() { |
| 21 | + It("reports an explicit error instead of handing out an unbindable port", func() { |
| 22 | + // The allocator used to increment without an upper bound. Past |
| 23 | + // 65535 it kept returning larger integers, which cannot be bound, |
| 24 | + // so the operator saw "backend won't start" with nothing pointing |
| 25 | + // at the allocator. Exhaustion must be named as exhaustion. |
| 26 | + s := &backendSupervisor{ |
| 27 | + processes: map[string]*backendProcess{}, |
| 28 | + minPort: 50051, |
| 29 | + nextPort: 50051, |
| 30 | + maxPort: 50052, |
| 31 | + portQuarantine: time.Hour, |
| 32 | + } |
| 33 | + |
| 34 | + first := mustAllocate(s, "a#0") |
| 35 | + second := mustAllocate(s, "b#0") |
| 36 | + Expect([]int{first, second}).To(ConsistOf(50051, 50052)) |
| 37 | + |
| 38 | + _, err := s.allocatePort("c#0") |
| 39 | + Expect(err).To(MatchError(ErrNoFreePort)) |
| 40 | + // The operator has to learn which knob to turn from the message |
| 41 | + // alone; a bare "no free port" sends them reading source. |
| 42 | + Expect(err.Error()).To(ContainSubstring("50051")) |
| 43 | + Expect(err.Error()).To(ContainSubstring("50052")) |
| 44 | + Expect(err.Error()).To(ContainSubstring("LOCALAI_GRPC_MAX_PORT")) |
| 45 | + }) |
| 46 | + |
| 47 | + It("does not hand out a port beyond the end of the range", func() { |
| 48 | + s := &backendSupervisor{ |
| 49 | + processes: map[string]*backendProcess{}, |
| 50 | + minPort: 50051, |
| 51 | + nextPort: 50051, |
| 52 | + maxPort: 50053, |
| 53 | + portQuarantine: time.Hour, |
| 54 | + } |
| 55 | + |
| 56 | + for i := 0; i < 3; i++ { |
| 57 | + port := mustAllocate(s, "model#0") |
| 58 | + Expect(port).To(BeNumerically(">=", 50051)) |
| 59 | + Expect(port).To(BeNumerically("<=", 50053)) |
| 60 | + } |
| 61 | + }) |
| 62 | + }) |
| 63 | + |
| 64 | + Describe("unexpectedly dead processes", func() { |
| 65 | + It("returns the dead process's port to the allocator", func() { |
| 66 | + // The death path deleted the map entry without releasing the port, |
| 67 | + // so every unexpected exit permanently consumed one port. A |
| 68 | + // crash-looping backend leaks one per restart, which is the real |
| 69 | + // route to exhausting the range — not the concurrent-peak growth |
| 70 | + // the issue assumed. |
| 71 | + bp := &backendProcess{port: 50051} |
| 72 | + s := &backendSupervisor{ |
| 73 | + processes: map[string]*backendProcess{"model#0": bp}, |
| 74 | + minPort: 50051, |
| 75 | + nextPort: 50052, |
| 76 | + maxPort: 50060, |
| 77 | + portQuarantine: time.Millisecond, |
| 78 | + } |
| 79 | + |
| 80 | + s.reapDeadProcess("model#0", bp) |
| 81 | + |
| 82 | + Expect(s.processes).NotTo(HaveKey("model#0")) |
| 83 | + // Quarantined, not lost: the port must come back. |
| 84 | + Expect(quarantinedPortNumbers(s)).To(ContainElement(50051)) |
| 85 | + time.Sleep(5 * time.Millisecond) |
| 86 | + Expect(mustAllocate(s, "model#0")).To(Equal(50051)) |
| 87 | + }) |
| 88 | + |
| 89 | + It("does not exhaust the range when one backend crash-loops", func() { |
| 90 | + // Restarting the same key repeatedly must not walk the allocator up |
| 91 | + // the range. Before the fix this consumed a fresh port per restart. |
| 92 | + s := &backendSupervisor{ |
| 93 | + processes: map[string]*backendProcess{}, |
| 94 | + minPort: 50051, |
| 95 | + nextPort: 50051, |
| 96 | + maxPort: 50053, |
| 97 | + portQuarantine: time.Millisecond, |
| 98 | + } |
| 99 | + |
| 100 | + for i := 0; i < 10; i++ { |
| 101 | + port, err := s.allocatePort("crashy#0") |
| 102 | + Expect(err).NotTo(HaveOccurred(), "restart %d exhausted the range", i) |
| 103 | + bp := &backendProcess{port: port} |
| 104 | + s.processes["crashy#0"] = bp |
| 105 | + s.reapDeadProcess("crashy#0", bp) |
| 106 | + time.Sleep(2 * time.Millisecond) |
| 107 | + } |
| 108 | + }) |
| 109 | + }) |
| 110 | + |
| 111 | + Describe("effectiveMaxPort", func() { |
| 112 | + It("uses the full range when unset", func() { |
| 113 | + Expect((&Config{}).effectiveMaxPort(50051)).To(Equal(65535)) |
| 114 | + }) |
| 115 | + |
| 116 | + It("honours an operator-set ceiling", func() { |
| 117 | + Expect((&Config{GRPCMaxPort: 51000}).effectiveMaxPort(50051)).To(Equal(51000)) |
| 118 | + }) |
| 119 | + |
| 120 | + It("clamps a ceiling above the highest TCP port", func() { |
| 121 | + Expect((&Config{GRPCMaxPort: 70000}).effectiveMaxPort(50051)).To(Equal(65535)) |
| 122 | + }) |
| 123 | + |
| 124 | + It("ignores a ceiling below the base port instead of wedging every start", func() { |
| 125 | + // An inverted range would leave the allocator with nothing to hand |
| 126 | + // out, so a typo here would take the whole worker down rather than |
| 127 | + // degrade it. |
| 128 | + Expect((&Config{GRPCMaxPort: 40000}).effectiveMaxPort(50051)).To(Equal(65535)) |
| 129 | + }) |
| 130 | + }) |
| 131 | + |
| 132 | + Describe("per-key port affinity", func() { |
| 133 | + It("never hands one key's released port to a different key", func() { |
| 134 | + // This is what makes a stale NodeModel row harmless. Process keys |
| 135 | + // (modelID#replica) and NodeModel rows (nodeID, modelName, |
| 136 | + // replicaIndex) are isomorphic, so if a port can only ever be |
| 137 | + // re-bound by the key that last held it, the only row that can name |
| 138 | + // that port is that key's own row — which its re-registration |
| 139 | + // overwrites. Misrouting to a *different* model becomes impossible |
| 140 | + // by construction rather than by racing a quarantine timer. |
| 141 | + s := &backendSupervisor{ |
| 142 | + processes: map[string]*backendProcess{}, |
| 143 | + minPort: 50051, |
| 144 | + nextPort: 50051, |
| 145 | + maxPort: 50060, |
| 146 | + portQuarantine: time.Millisecond, |
| 147 | + } |
| 148 | + |
| 149 | + portA := mustAllocate(s, "alpha#0") |
| 150 | + s.releasePortForKey("alpha#0", portA) |
| 151 | + time.Sleep(5 * time.Millisecond) |
| 152 | + |
| 153 | + // beta must NOT be given alpha's port while unused ports remain. |
| 154 | + portB := mustAllocate(s, "beta#0") |
| 155 | + Expect(portB).NotTo(Equal(portA)) |
| 156 | + |
| 157 | + // alpha coming back gets its own port again. |
| 158 | + Expect(mustAllocate(s, "alpha#0")).To(Equal(portA)) |
| 159 | + }) |
| 160 | + |
| 161 | + It("reuses an unowned port before growing the range", func() { |
| 162 | + s := &backendSupervisor{ |
| 163 | + processes: map[string]*backendProcess{}, |
| 164 | + minPort: 50051, |
| 165 | + nextPort: 50051, |
| 166 | + maxPort: 50060, |
| 167 | + portQuarantine: time.Millisecond, |
| 168 | + } |
| 169 | + |
| 170 | + // A port released without an owning key (legacy/unknown) is free |
| 171 | + // for anyone: no row can be attributed to it. |
| 172 | + s.releasePort(50051) |
| 173 | + time.Sleep(5 * time.Millisecond) |
| 174 | + |
| 175 | + Expect(mustAllocate(s, "gamma#0")).To(Equal(50051)) |
| 176 | + }) |
| 177 | + |
| 178 | + It("steals another key's port rather than failing an exhausted range", func() { |
| 179 | + // Affinity is a preference, not a reservation. Refusing to start a |
| 180 | + // backend because a long-gone key still owns the last port trades a |
| 181 | + // vanishingly rare misroute for a guaranteed outage. |
| 182 | + s := &backendSupervisor{ |
| 183 | + processes: map[string]*backendProcess{}, |
| 184 | + minPort: 50051, |
| 185 | + nextPort: 50051, |
| 186 | + maxPort: 50051, |
| 187 | + portQuarantine: time.Millisecond, |
| 188 | + } |
| 189 | + |
| 190 | + portA := mustAllocate(s, "alpha#0") |
| 191 | + s.releasePortForKey("alpha#0", portA) |
| 192 | + time.Sleep(5 * time.Millisecond) |
| 193 | + |
| 194 | + port, err := s.allocatePort("beta#0") |
| 195 | + Expect(err).NotTo(HaveOccurred()) |
| 196 | + Expect(port).To(Equal(portA)) |
| 197 | + }) |
| 198 | + |
| 199 | + It("stops sequestering the port once the stale-row window has passed", func() { |
| 200 | + // Affinity exists to outlive any controller row that could still |
| 201 | + // name the port. Past that, holding it for a key that may never |
| 202 | + // come back is pure cost: every distinct model the worker has ever |
| 203 | + // served would consume a port permanently, so a long-lived worker |
| 204 | + // would climb to the end of its range on distinct-key count rather |
| 205 | + // than on concurrency, then steal on every subsequent allocation |
| 206 | + // while advising the operator to raise a ceiling that is not the |
| 207 | + // problem. |
| 208 | + s := &backendSupervisor{ |
| 209 | + processes: map[string]*backendProcess{}, |
| 210 | + minPort: 50051, |
| 211 | + nextPort: 50051, |
| 212 | + maxPort: 50060, |
| 213 | + portQuarantine: time.Millisecond, |
| 214 | + portAffinityWindow: 20 * time.Millisecond, |
| 215 | + } |
| 216 | + |
| 217 | + portA := mustAllocate(s, "alpha#0") |
| 218 | + s.releasePortForKey("alpha#0", portA) |
| 219 | + time.Sleep(60 * time.Millisecond) |
| 220 | + |
| 221 | + // A brand-new key reuses it rather than growing the range. |
| 222 | + Expect(mustAllocate(s, "beta#0")).To(Equal(portA)) |
| 223 | + }) |
| 224 | + |
| 225 | + It("applies the default affinity window when none is configured", func() { |
| 226 | + // The zero value must not degrade to "no affinity": every |
| 227 | + // supervisor built outside the tests leaves the field unset, and |
| 228 | + // zero would hand a just-released port straight to another model. |
| 229 | + s := &backendSupervisor{ |
| 230 | + processes: map[string]*backendProcess{}, |
| 231 | + minPort: 50051, |
| 232 | + nextPort: 50051, |
| 233 | + maxPort: 50060, |
| 234 | + portQuarantine: time.Millisecond, |
| 235 | + } |
| 236 | + |
| 237 | + portA := mustAllocate(s, "alpha#0") |
| 238 | + s.releasePortForKey("alpha#0", portA) |
| 239 | + time.Sleep(5 * time.Millisecond) |
| 240 | + |
| 241 | + Expect(mustAllocate(s, "beta#0")).NotTo(Equal(portA)) |
| 242 | + }) |
| 243 | + |
| 244 | + It("keeps the affinity map bounded by the size of the port range", func() { |
| 245 | + // The affinity map must not become the port leak's twin. Handing a |
| 246 | + // port to a new key drops the previous owner's entry, so the map is |
| 247 | + // injective over ports and can never hold more entries than the |
| 248 | + // range has ports, however many distinct keys the worker sees. |
| 249 | + s := &backendSupervisor{ |
| 250 | + processes: map[string]*backendProcess{}, |
| 251 | + minPort: 50051, |
| 252 | + nextPort: 50051, |
| 253 | + maxPort: 50052, |
| 254 | + portQuarantine: time.Nanosecond, |
| 255 | + } |
| 256 | + |
| 257 | + for i := 0; i < 50; i++ { |
| 258 | + key := "model" + string(rune('a'+i%26)) + "#0" |
| 259 | + port, err := s.allocatePort(key) |
| 260 | + Expect(err).NotTo(HaveOccurred()) |
| 261 | + s.releasePortForKey(key, port) |
| 262 | + time.Sleep(time.Millisecond) |
| 263 | + } |
| 264 | + |
| 265 | + Expect(len(s.portAffinity)).To(BeNumerically("<=", 2)) |
| 266 | + }) |
| 267 | + }) |
| 268 | +}) |
0 commit comments