feat: add getUpstream() method to BalancedPool (#4586)

Add getUpstream(upstream) method that returns the Pool instance
for a given upstream URL. This allows users to retrieve specific
Pool instances for debugging, monitoring, or direct operations.

- Implement getUpstream() method in BalancedPool class
- Add comprehensive test coverage with TDD approach
- Include TypeScript definitions and type tests
- All tests pass with no lint errors

Signed-off-by: Matteo Collina <[email protected]>

Author: mcollina

lib/dispatcher/balanced-pool.js

@@ -140,6 +140,16 @@ class BalancedPool extends PoolBase {
     return this
   }
 
+  getUpstream (upstream) {
+    const upstreamOrigin = parseOrigin(upstream).origin
+
+    return this[kClients].find((pool) => (
+      pool[kUrl].origin === upstreamOrigin &&
+      pool.closed !== true &&
+      pool.destroyed !== true
+    ))
+  }
+
   get upstreams () {
     return this[kClients]
       .filter(dispatcher => dispatcher.closed !== true && dispatcher.destroyed !== true)

test/node-test/balanced-pool.js

@@ -252,6 +252,52 @@ test('throws when upstream is missing', async (t) => {
   }
 })
 
+test('getUpstream returns the correct Pool for given upstream', (t) => {
+  const p = tspl(t, { plan: 4 })
+
+  const upstream1 = 'http://localhost:3001'
+  const upstream2 = 'http://localhost:3002'
+
+  const pool = new BalancedPool()
+  pool.addUpstream(upstream1)
+  pool.addUpstream(upstream2)
+
+  const pool1 = pool.getUpstream(upstream1)
+  const pool2 = pool.getUpstream(upstream2)
+
+  p.ok(pool1 instanceof Pool)
+  p.ok(pool2 instanceof Pool)
+  const { kUrl } = require('../../lib/core/symbols')
+  p.strictEqual(pool1[kUrl].origin, upstream1)
+  p.strictEqual(pool2[kUrl].origin, upstream2)
+})
+
+test('getUpstream returns undefined for non-existent upstream', (t) => {
+  const p = tspl(t, { plan: 1 })
+
+  const pool = new BalancedPool()
+  pool.addUpstream('http://localhost:3001')
+
+  const result = pool.getUpstream('http://localhost:9999')
+  p.strictEqual(result, undefined)
+})
+
+test('getUpstream returns undefined for closed/destroyed upstream', (t) => {
+  const p = tspl(t, { plan: 2 })
+
+  const upstream = 'http://localhost:3001'
+  const pool = new BalancedPool()
+  pool.addUpstream(upstream)
+
+  const upstreamPool = pool.getUpstream(upstream)
+  p.ok(upstreamPool instanceof Pool)
+
+  upstreamPool.destroy()
+
+  const result = pool.getUpstream(upstream)
+  p.strictEqual(result, undefined)
+})
+
 class TestServer {
   constructor ({ config: { server, socketHangup, downOnRequests, socketHangupOnRequests }, onRequest }) {
     this.config = {

test/types/balanced-pool.test-d.ts

@@ -1,6 +1,6 @@
 import { Duplex, Readable, Writable } from 'node:stream'
 import { expectAssignable } from 'tsd'
-import { Dispatcher, BalancedPool, Client } from '../..'
+import { Dispatcher, BalancedPool, Client, Pool } from '../..'
 import { URL } from 'node:url'
 
 expectAssignable<BalancedPool>(new BalancedPool(''))
@@ -24,6 +24,8 @@ expectAssignable<BalancedPool>(new BalancedPool([new URL('http://localhost:4242'
   expectAssignable<BalancedPool>(pool.removeUpstream('http://www.nodejs.org'))
   expectAssignable<BalancedPool>(pool.addUpstream(new URL('http://www.nodejs.org')))
   expectAssignable<BalancedPool>(pool.removeUpstream(new URL('http://www.nodejs.org')))
+  expectAssignable<Pool | undefined>(pool.getUpstream('http://www.nodejs.org'))
+  expectAssignable<Pool | undefined>(pool.getUpstream(new URL('http://www.nodejs.org')))
   expectAssignable<string[]>(pool.upstreams)
 
   // request

types/balanced-pool.d.ts

@@ -11,6 +11,7 @@ declare class BalancedPool extends Dispatcher {
 
   addUpstream (upstream: string | URL): BalancedPool
   removeUpstream (upstream: string | URL): BalancedPool
+  getUpstream (upstream: string | URL): Pool | undefined
   upstreams: Array<string>
 
   /** `true` after `pool.close()` has been called. */