elevate-adapter: connection pooling to superpowerd
Reuse SEQPACKET connections across requests instead of connect/close per command. Connections are returned to a pool (max 4 idle) after successful use and discarded on error. superpowerd HandleConn already loops for multiple requests per connection, so this is safe. Each ExecuteCommandStreaming still blocks its connection until the command exits, so concurrent commands use separate pooled connections.
This commit is contained in:
parent
0fa24ef8c2
commit
2803fb30f3
|
|
@ -35,6 +35,68 @@ _token_lock = threading.Lock()
|
|||
_LOG_PREFIX = "elevate-adapter-superpowerd"
|
||||
|
||||
|
||||
class _SpPool:
|
||||
"""Pool of pre-connected SEQPACKET sockets to superpowerd.
|
||||
|
||||
Connections are created on demand and returned after use.
|
||||
Failed connections are discarded, not returned.
|
||||
"""
|
||||
|
||||
def __init__(self, max_idle=4):
|
||||
self._lock = threading.Lock()
|
||||
self._idle = [] # list of ready-to-use sockets
|
||||
self._max_idle = max_idle
|
||||
|
||||
def acquire(self):
|
||||
"""Get a connected socket, from pool or freshly created."""
|
||||
with self._lock:
|
||||
while self._idle:
|
||||
sp = self._idle.pop()
|
||||
# Verify connection is still alive (peek for hangup)
|
||||
try:
|
||||
sp.setblocking(False)
|
||||
try:
|
||||
# If recv returns 0 or raises, connection is dead
|
||||
data = sp.recv(1, socket.MSG_PEEK | socket.MSG_DONTWAIT)
|
||||
if not data:
|
||||
sp.close()
|
||||
continue
|
||||
except BlockingIOError:
|
||||
# No data available = good, connection alive
|
||||
pass
|
||||
except OSError:
|
||||
sp.close()
|
||||
continue
|
||||
finally:
|
||||
sp.setblocking(True)
|
||||
return sp
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
# Pool empty — create a new connection
|
||||
sp = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
|
||||
sp.connect(sp_socket_path())
|
||||
return sp
|
||||
|
||||
def release(self, sp):
|
||||
"""Return a healthy connection to the pool."""
|
||||
with self._lock:
|
||||
if len(self._idle) < self._max_idle:
|
||||
self._idle.append(sp)
|
||||
else:
|
||||
sp.close()
|
||||
|
||||
def close_all(self):
|
||||
"""Close all pooled connections."""
|
||||
with self._lock:
|
||||
for sp in self._idle:
|
||||
sp.close()
|
||||
self._idle.clear()
|
||||
|
||||
|
||||
_sp_pool = _SpPool()
|
||||
|
||||
|
||||
def _log(msg):
|
||||
print(f"{_LOG_PREFIX}: {msg}", file=sys.stderr, flush=True)
|
||||
|
||||
|
|
@ -177,14 +239,11 @@ def _is_suspicious_failure(sp_data, output_received):
|
|||
def _execute_on_superpowerd(conn, cmd, cwd, token):
|
||||
"""Execute a command via superpowerd. Returns (sp_data, True) on success/normal
|
||||
error, or (sp_data, False) if the caller should retry (auth error, no output streamed)."""
|
||||
sp_path = sp_socket_path()
|
||||
|
||||
# Connect to superpowerd via SEQPACKET
|
||||
sp = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
|
||||
# Acquire a connection from the pool
|
||||
try:
|
||||
sp.connect(sp_path)
|
||||
sp = _sp_pool.acquire()
|
||||
except OSError as e:
|
||||
send_error(conn, f"elevate-adapter: connect {sp_path}: {e}")
|
||||
send_error(conn, f"elevate-adapter: connect {sp_socket_path()}: {e}")
|
||||
return None, True
|
||||
|
||||
# Create pipes for stdout+stderr (merged) and stdin (/dev/null)
|
||||
|
|
@ -244,11 +303,18 @@ def _execute_on_superpowerd(conn, cmd, cwd, token):
|
|||
t.start()
|
||||
|
||||
# Wait for superpowerd's response (comes after command exits)
|
||||
sp_ok = True
|
||||
try:
|
||||
sp_data = sp.recv(1024 * 1024)
|
||||
except OSError:
|
||||
sp_data = b""
|
||||
sp.close()
|
||||
sp_ok = False
|
||||
|
||||
# Return connection to pool if healthy, otherwise discard
|
||||
if sp_ok and sp_data:
|
||||
_sp_pool.release(sp)
|
||||
else:
|
||||
sp.close()
|
||||
|
||||
# Wait for streaming thread to finish
|
||||
stream_done.wait()
|
||||
|
|
|
|||
Loading…
Reference in New Issue