Package biopb.tensor

Class TensorConnectionPool

java.lang.Object

biopb.tensor.TensorConnectionPool

public final class TensorConnectionPool
extends Object

Thread-local connection pool for TensorFlightClient connections. This pool is used by SerializableTensorImg to reconstruct connections after deserialization. Each thread maintains its own FlightClient for lock-free read access, with automatic cleanup on thread death or JVM shutdown. Design mirrors Python's thread-local connection pool with: - Per-thread FlightClient storage (lock-free reads) - Connection registry for cleanup - Thread death eviction - Shutdown hook for resource cleanup

Nested Class Summary

Nested Classes

Modifier and Type	Class	Description
static final class	TensorConnectionPool.PooledConnection	A pooled connection containing FlightClient and authentication option.
static final class	TensorConnectionPool.SharedCache	Shared cache for cross-thread cache hits.

Method Summary

Modifier and Type	Method	Description
static void	clear()	Clear all pooled connections (for testing).
static void	closeConnection(org.apache.arrow.flight.Location location, String token)	Close a specific connection for the given location and token.
static TensorConnectionPool.PooledConnection	getConnection(org.apache.arrow.flight.Location location, String token, org.apache.arrow.memory.BufferAllocator allocator)	Get or create a pooled connection for the given location and token.
static int	getConnectionCount()	Get the total connection count across all threads (for testing/monitoring).
static TensorConnectionPool.SharedCache	getSharedCache(biopb.tensor.TensorConnectionPool.ConnectionKey key, long maxBytes)	Get or create a shared cache for the given connection key.
static int	getThreadCount()	Get the thread count with active connections (for testing/monitoring).

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Method Details

getConnection

public static TensorConnectionPool.PooledConnection getConnection(org.apache.arrow.flight.Location location,
 String token,
 org.apache.arrow.memory.BufferAllocator allocator)

Get or create a pooled connection for the given location and token. Uses thread-local storage for lock-free reads. New connections are registered for cleanup when the thread dies or JVM shuts down.

Parameters:

location - Flight server location

token - Bearer token for authentication (null disables auth)

allocator - Buffer allocator for the connection

Returns:

PooledConnection containing FlightClient and auth option

closeConnection

public static void closeConnection(org.apache.arrow.flight.Location location,
 String token)

Close a specific connection for the given location and token. Removes the connection from the current thread's pool and the registry.

Parameters:

location - Flight server location

token - Bearer token for authentication

getSharedCache

public static TensorConnectionPool.SharedCache getSharedCache(biopb.tensor.TensorConnectionPool.ConnectionKey key,
 long maxBytes)

Get or create a shared cache for the given connection key. Shared caches allow cross-thread cache hits for read-heavy workloads.

Parameters:

key - Connection key

maxBytes - Maximum bytes to cache

Returns:

SharedCache for the connection

getConnectionCount

public static int getConnectionCount()

Get the total connection count across all threads (for testing/monitoring).

getThreadCount

public static int getThreadCount()

Get the thread count with active connections (for testing/monitoring).

clear

public static void clear()

Clear all pooled connections (for testing).