Class: Neo4j::Driver::Routing::LoadBalancer

Inherits:
Object
  • Object
show all
Defined in:
lib/neo4j/driver/routing/load_balancer.rb

Overview

Routing-aware connection provider. Mirrors the design of Java's RoutingTableHandlerImpl + LoadBalancer and Python's Neo4jPool:

  • Per-database RoutingTable cache, mutated through @refresh_lock.
  • Per-server connection pools (TimedStack), keyed by ServerAddress.
  • acquire(access_mode:) ensures the table is fresh, then loops selecting an address and opening a connection; on connection failure the address is deactivated (forgotten from every table + per-server pool torn down) and we retry until the role bucket is exhausted.
  • Connections handed out are wrapped in RoutedConnection so the pool gets called back on write-side leader changes (NotALeader, ForbiddenOnReadOnlyDatabase) and stale-server connection errors.

Defined Under Namespace

Classes: Handler

Constant Summary collapse

ROUTING_CONTEXT_RESERVED_KEYS =
%w[address].freeze
FATAL_DISCOVERY_CODES =

Codes that mean "this is a client mistake, retrying the routing fetch against another router won't help, fail fast". Matches Python's Neo4jError._is_fatal_during_discovery.

%w[
  Neo.ClientError.Database.DatabaseNotFound
  Neo.ClientError.Transaction.InvalidBookmark
  Neo.ClientError.Transaction.InvalidBookmarkMixture
  Neo.ClientError.Statement.TypeError
  Neo.ClientError.Statement.ArgumentError
  Neo.ClientError.Request.Invalid
].freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(uri, auth_manager, options = {}, domain_name_resolver: nil, clock: Internal::Clock.new) ⇒ LoadBalancer

domain_name_resolver is the factory-injected hostname->IPs hook (nil = system DNS); baked into every connection this balancer builds.



36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 36

def initialize(uri, auth_manager, options = {}, domain_name_resolver: nil, clock: Internal::Clock.new)
  @uri = uri
  @auth_manager = auth_manager
  @options = options
  @clock = clock
  @domain_name_resolver = domain_name_resolver
  @routing_context = parse_routing_context(uri)
  @pools = {} # ServerAddress => ConnectionPool::TimedStack
  @routing_tables = {} # database (str or nil) => RoutingTable
  @cursor = Hash.new(0) # round-robin per (database, role)
  # Per-server authorization-expired generation counters (see
  # Connection#auth_epoch and Direct::ConnectionProvider). Bumped for a
  # server when one of its connections reports AuthorizationExpired, so
  # the OTHER connections to that same server re-authenticate on their
  # next acquire. Scoped per-address (not driver-wide) because the
  # server's authorization cache is per-server — a reader's expiry must
  # not force the writer pool to re-auth.
  @auth_epochs = Hash.new(0)
  # Monitor (reentrant) because ensure_routing_table_is_fresh holds
  # the lock while it goes through pool_for, which also locks.
  @refresh_lock = Monitor.new
  # Optimization:HomeDatabaseCache — a driver-wide identity->home-db
  # cache, consulted only when every open pooled connection advertises
  # server-side routing (@ssr_with/@ssr_without track that tally, kept
  # current by track_ssr_open/close as connections open and tear down).
  @home_db_cache = Internal::HomeDbCache.new
  @ssr_with = 0
  @ssr_without = 0
  @ssr_lock = Mutex.new
end

Instance Attribute Details

#home_db_cacheObject (readonly)

Returns the value of attribute home_db_cache.



67
68
69
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 67

def home_db_cache
  @home_db_cache
end

Instance Method Details

#acquire(access_mode: :write, database: nil, bookmarks: nil, imp_user: nil, auth: nil, deadline: nil) ⇒ Object

Open (or pop) a connection appropriate for access_mode against database (nil = home db). Loops: select an address, try to acquire, deactivate on connection failure and try again. Raises ServiceUnavailableException only when the role bucket has been exhausted by deactivations.



93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 93

def acquire(access_mode: :write, database: nil, bookmarks: nil, imp_user: nil, auth: nil, deadline: nil)
  # Fast-fail with IllegalStateException for use-after-close.
  # Otherwise the next routing fetch would propagate a generic
  # Connection-refused ServiceUnavailableException, masking the
  # actual bug in the caller's lifecycle.
  raise Exceptions::IllegalStateException, 'Driver is closed' if @closed

  access_mode = access_mode.to_sym
  # `auth` is the per-session override (nil = manager's current
  # token). The worker identity is resolved per turn inside the loop
  # below, not once here: on Bolt 5.0 an acquire that discards a
  # token-rotated connection must re-consult the manager so the
  # replacement issues its own get_token (mirrors
  # Direct::ConnectionProvider — one extra get_auth per rotation).
  # Discovery resolves its own identity independently, so a routed
  # acquire with the default token consults the manager twice
  # (get_auth_count == 2): once for the ROUTE connection, once for the
  # worker. A session-carried token short-circuits both (count 0).
  #
  # imp_user is threaded into discovery so the ROUTE call enforces
  # impersonation support (Bolt 4.4+) against the router, matching
  # the RUN/BEGIN path — see Connection#route. `auth` is threaded so
  # the ROUTE connection authenticates as the session's identity
  # (per-session token) rather than always the manager's default.
  # Resolve to the concrete database name: for a home-db acquire
  # (database == nil) discovery returns the resolved name, which keys
  # the table and the address selection below.
  resolved_database = ensure_routing_table_is_fresh(
    database, access_mode, bookmarks: bookmarks, imp_user: imp_user, auth: auth
  ).database

  last_error = nil
  loop do
    address = select_address(resolved_database, access_mode)
    unless address
      # Routing table yielded no usable server for this mode —
      # a session-expired condition (the session can't be served,
      # caller should get a fresh one), even if the last attempt
      # failed with a connection-level ServiceUnavailable. Surface
      # SessionExpired with that as the cause (matches Java).
      # Explicit cause: this raise is outside the per-address
      # rescue, so $! would not auto-populate it. cause: nil is
      # fine when there was no connection-level failure.
      raise Exceptions::SessionExpiredException,
            "No #{access_mode} servers available for database #{resolved_database.inspect}",
            cause: last_error
    end

    pool = pool_for(address)
    begin
      epoch = auth_epoch_for(address)
      # Re-resolve per turn (see acquire header): a discard-and-retry
      # on Bolt 5.0 token rotation must issue its own get_token.
      effective = auth || @auth_manager.get_token
      inner = pool.pop(auth: effective, deadline: deadline)
      begin
        ensure_identity(inner, effective, session_auth: auth, address: address, epoch: epoch)
      rescue StandardError
        # Don't leak the worker slot if identity enforcement fails
        # (per-session auth on Bolt < 5.1, or a re-auth LOGON
        # failure); discard it, then let the error propagate.
        discard(address, inner)
        raise
      end
      # On Bolt 5.0 ensure_identity can't re-auth in place, so a
      # reused connection whose token or auth epoch is stale (token
      # rotation, or an AuthorizationExpired refresh) must be discarded
      # and replaced by a fresh one — mirrors Direct::ConnectionProvider.
      unless inner.protocol.supports_re_auth? ||
             (inner.auth == effective && inner.auth_epoch == epoch)
        discard(address, inner)
        next
      end
      return RoutedConnection.new(self, inner, address, access_mode, resolved_database)
    rescue Exceptions::ServiceUnavailableException => e
      # Server is unreachable (open_connection raised inside the
      # pool's create block). Drop the address from every table
      # and tear down its pool; loop and try another address.
      last_error = e
      deactivate(address)
    end
  end
end

#auth_epoch_for(address) ⇒ Object

Current auth generation for a server, read under the lock that guards the @auth_epochs Hash (mutated in on_security_exception).



250
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 250

def auth_epoch_for(address) = @refresh_lock.synchronize { @auth_epochs[address] }

#cache_home_db(imp_user, auth, database) ⇒ Object

Record the server's resolved home database for this identity.



217
218
219
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 217

def cache_home_db(imp_user, auth, database)
  @home_db_cache.set(@home_db_cache.compute_key(imp_user, auth), database)
end

#closeObject



319
320
321
322
323
324
325
326
327
328
329
330
331
332
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 319

def close
  @refresh_lock.synchronize do
    @closed = true
    @pools.each_value do |pool|
      pool.shutdown do |conn|
        conn.close
      rescue StandardError
        nil
      end
    end
    @pools.clear
    @routing_tables.clear
  end
end

#connection_pool_metricsObject

Per-server-address pool snapshots (driver.metrics). One entry per server we hold a pool for; the address key is the Routing::ServerAddress the pool was created under.



293
294
295
296
297
298
299
300
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 293

def connection_pool_metrics
  @refresh_lock.synchronize do
    @pools.map do |address, pool|
      in_use, idle = pool.metrics_snapshot
      Internal::Metrics::ConnectionPoolMetrics.new(address, in_use, idle)
    end
  end
end

#current_auth_tokenObject



232
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 232

def current_auth_token = @auth_manager.get_token

#deactivate(address) ⇒ Object

Called by RoutedConnection on a fatal connection-level error (or DatabaseUnavailable). Removes the address from every database's routing table and tears down its connection pool.



376
377
378
379
380
381
382
383
384
385
386
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 376

def deactivate(address)
  @refresh_lock.synchronize do
    @routing_tables.each_value { |table| table.forget(address) }
    pool = @pools.delete(address)
    pool&.shutdown do |conn|
      conn.close
    rescue StandardError
      nil
    end
  end
end

#home_database(bookmarks, imp_user = nil, auth = nil) ⇒ Object

Current default identity from the auth-token manager (refreshes as needed); on_security_exception feeds a failure back to it. Mirror Direct::ConnectionProvider so Session stays polymorphic. Resolve the user's home database (database == nil): the ROUTE response carries the resolved name in db (Bolt 4.4+/5.x), which the routing table records as its database. The session uses the resolved name on RUN/BEGIN so the server doesn't re-resolve it per op. nil on the procedure path (3.0/4.0-4.2 have no db in the reply), where the server resolves the home db from a null db itself.



186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 186

def home_database(bookmarks, imp_user = nil, auth = nil)
  # Same use-after-close guard as #acquire: home-db resolution routes
  # too, so a closed driver must fast-fail with IllegalStateException
  # rather than re-routing into a Connection-refused
  # ServiceUnavailableException. This is the ONLY routing entry for the
  # single-database path (Bolt 3.0, database == nil) — it never reaches
  # #acquire's guard — so without this a post-close session.run on a 3.0
  # routing driver surfaced the wrong error type.
  raise Exceptions::IllegalStateException, 'Driver is closed' if @closed

  resolved = ensure_routing_table_is_fresh(nil, :read, bookmarks: bookmarks, imp_user: imp_user,
                                                       auth: auth).database
  # Remember the authoritative name (Optimization:HomeDatabaseCache) so a
  # later same-identity session can guess it and skip discovery.
  cache_home_db(imp_user, auth, resolved)
  resolved
end

#home_db_guess(imp_user, auth) ⇒ Object

Optimization:HomeDatabaseCache — this identity's last resolved home database, or nil when we can't optimistically guess it: no cache entry, or not every open connection does server-side routing (so the server might not re-route a stale guess). Used only to pick the routing table to acquire against — the operation still sends db=null so the server resolves the real home db.



210
211
212
213
214
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 210

def home_db_guess(imp_user, auth)
  return nil unless ssr_enabled?

  @home_db_cache.get(@home_db_cache.compute_key(imp_user, auth))
end

#on_security_exception(address, token, error, session_scoped = false) ⇒ Object



234
235
236
237
238
239
240
241
242
243
244
245
246
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 234

def on_security_exception(address, token, error, session_scoped = false)
  # See Direct::ConnectionProvider#on_security_exception: an expired
  # authorization cache forces the OTHER connections to that same
  # server to re-auth. Provider-side, so it runs regardless of who
  # owns the token; the manager notification is skipped for a
  # per-session identity.
  if error.is_a?(Exceptions::AuthorizationExpiredException)
    @refresh_lock.synchronize { @auth_epochs[address] += 1 }
  end
  return false if session_scoped

  @auth_manager.handle_security_exception(token, error)
end

#on_write_failure(address, database) ⇒ Object

Called by RoutedConnection on a write-mode operation that hit NotALeader / ForbiddenOnReadOnlyDatabase. The server is alive but no longer the leader for this db — drop it from the writers bucket only; routers/readers stay.



392
393
394
395
396
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 392

def on_write_failure(address, database)
  @refresh_lock.synchronize do
    @routing_tables[database]&.forget_writer(address)
  end
end

#refresh(database, bookmarks = nil) ⇒ Object

Internal — force a fresh ROUTE call for the given database regardless of TTL/cache. Used by ForcedRoutingTableUpdate testkit handler. Threads bookmarks through to the ROUTE payload so causal-consistency assertions work.



362
363
364
365
366
367
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 362

def refresh(database, bookmarks = nil)
  @refresh_lock.synchronize do
    invalidate_routing_table(database)
    update_routing_table(database, bookmarks: bookmarks)
  end
end

#release(connection) ⇒ Object



252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 252

def release(connection)
  # Direct provider tolerates nil; mirror that. Internal callers
  # always release a RoutedConnection but guard so the wrong
  # type doesn't NoMethodError later.
  return unless connection.respond_to?(:discard_on_release)

  if connection.discard_on_release
    discard(connection.address_obj, connection.inner)
    return
  end

  @refresh_lock.synchronize do
    pool = @pools[connection.address_obj]
    pool&.push(connection.inner)
  end
end

#routing_table_fresh?(database, access_mode = :read) ⇒ Boolean

Whether a usable table for database is already cached — so acquiring against it won't ROUTE. The home-db cache uses this to decide whether a guessed db can be sent as db=nil (table fresh, server resolves) or must be pinned (a ROUTE will run and authoritatively resolve it).

Returns:

  • (Boolean)


225
226
227
228
229
230
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 225

def routing_table_fresh?(database, access_mode = :read)
  @refresh_lock.synchronize do
    table = @routing_tables[database]
    !!table&.fresh?(readonly: access_mode == :read)
  end
end

#routing_table_handler(database) ⇒ Object

Internal — mirrors Java's RoutingTableRegistry#getRoutingTableHandler(databaseName). Pure read: returns the cached table for the database, or a new empty placeholder if no table has ever been fetched. (Empty means routers/readers/writers are all empty — it is NOT fresh?; callers that want a fetched table go through ensure_routing_table_is_fresh.) Deliberately does NOT force a fetch — testkit's get_routing_table contract is "what's currently known", and an auto-fetch here causes a second ROUTE on a stub server that already hung up after the first (see test_should_fail_on_routing_table_with_no_reader).



352
353
354
355
356
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 352

def routing_table_handler(database)
  @refresh_lock.synchronize do
    Handler.new(@routing_tables[database] || RoutingTable.new(database: database, clock: @clock))
  end
end

#routing_table_registryObject

Internal — mirrors Java's ConnectionProvider#getRoutingTableRegistry(). LoadBalancer is both the connection provider and the routing-table registry; the layer split exists in Java mostly because it predates generics. Used by testkit's GetRoutingTable handler.



339
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 339

def routing_table_registry = self

#ssr_enabled?Boolean

Home-db cache gate: true when there is at least one open pooled connection and every one advertised ssr.enabled, so the server will re-route an optimistically guessed home database.

Returns:

  • (Boolean)


72
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 72

def ssr_enabled? = @ssr_lock.synchronize { @ssr_with.positive? && @ssr_without.zero? }

#supports_multi_db?Boolean

Routing requires Bolt 4.0+ (the ROUTE message and CALL dbms.routing.getRoutingTable are 4.0+ features). If we got this far the answer is always true. Multi-database support is a property of the negotiated Bolt version (4.0+), not of routing per se — a neo4j:// driver against a 3.0 cluster still routes (via the getRoutingTable procedure) but does NOT support multiple databases. Probe an actual connection's protocol, mirroring Direct::ConnectionProvider. The acquire does discovery + a reader HELLO (no query), then the connection is released straight back to the pool.

Returns:

  • (Boolean)


312
313
314
315
316
317
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 312

def supports_multi_db?
  conn = acquire(access_mode: :read)
  conn.protocol.supports_multiple_databases?
ensure
  release(conn) if conn
end

#track_ssr_close(conn) ⇒ Object



78
79
80
81
82
83
84
85
86
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 78

def track_ssr_close(conn)
  @ssr_lock.synchronize do
    if conn.ssr_enabled?
      @ssr_with -= 1 if @ssr_with.positive?
    elsif @ssr_without.positive?
      @ssr_without -= 1
    end
  end
end

#track_ssr_open(conn) ⇒ Object



74
75
76
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 74

def track_ssr_open(conn)
  @ssr_lock.synchronize { conn.ssr_enabled? ? @ssr_with += 1 : @ssr_without += 1 }
end

#verify_connectivityObject

Routing-aware verify_connectivity: force-refresh the routing table, then probe any reader to confirm the cluster is reachable. RESET so the borrowed connection lands back in the pool in a known-clean state (matches testkit's test_routing_from_pool expectation of one RESET per probe).



274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
# File 'lib/neo4j/driver/routing/load_balancer.rb', line 274

def verify_connectivity
  invalidate_routing_table(nil)
  conn = acquire(access_mode: :read)
  begin
    # propagate: a failed probe means a dead reader — surface it and
    # discard the connection rather than pooling it (see
    # Direct::ConnectionProvider#verify_connectivity).
    conn.reset!(propagate: true)
  rescue StandardError
    conn.discard_on_release = true
    release(conn)
    raise
  end
  release(conn)
end