This commit is contained in:
2026-05-28 23:59:12 +08:00
parent 068d3f4be9
commit e80fe14624
4 changed files with 88 additions and 18 deletions
+60 -7
View File
@@ -19,6 +19,7 @@ final class AgentClient
{
private const MAX_SEND_BUFFER = 32 * 1024 * 1024;
private const BACKPRESSURE_HIGH_WATERMARK = 16 * 1024 * 1024;
private const DATA_CHUNK_SIZE = 256 * 1024;
private ?AsyncTcpConnection $pop = null;
private ?FrameParser $parser = null;
@@ -45,6 +46,10 @@ final class AgentClient
private array $pausedClientsForPop = [];
/** @var array<int, true> */
private array $clientsPausingPop = [];
/** @var array<string, true> */
private array $pendingClientCloses = [];
/** @var array<int, true> */
private array $suppressClientCloseFrames = [];
public function __construct(
private readonly string $clientListen,
@@ -132,6 +137,8 @@ final class AgentClient
$this->sessionIngressProtocols = [];
$this->pausedClientsForPop = [];
$this->clientsPausingPop = [];
$this->pendingClientCloses = [];
$this->suppressClientCloseFrames = [];
Timer::add(3, fn () => $this->connect(), [], false);
};
$connection->connect();
@@ -147,7 +154,7 @@ final class AgentClient
FrameType::OPEN_FAIL => $this->failClientSession($frame),
FrameType::DATA => $this->forwardToClient($frame),
FrameType::UDP_DATA => $this->forwardUdpToClient($frame),
FrameType::CLOSE => $this->closeClient($frame->sessionId),
FrameType::CLOSE => $this->closeClient($frame->sessionId, true),
default => null,
};
}
@@ -156,7 +163,7 @@ final class AgentClient
{
$connection->maxSendBufferSize = self::MAX_SEND_BUFFER;
$connection->onBufferFull = fn (TcpConnection $connection) => $this->pausePopForClient($connection);
$connection->onBufferDrain = fn (TcpConnection $connection) => $this->resumePopForClient($connection);
$connection->onBufferDrain = fn (TcpConnection $connection) => $this->onClientBufferDrain($connection);
$this->initialBuffers[$connection->id] = '';
$this->connectionStages[$connection->id] = 'init';
}
@@ -174,8 +181,7 @@ final class AgentClient
return;
}
$sent = $this->send(new Frame(FrameType::DATA, $sessionId, ['data_raw' => $data]));
if ($sent === false) {
if (!$this->sendData($sessionId, $data)) {
$this->closeClient($sessionId);
return;
}
@@ -564,7 +570,7 @@ final class AgentClient
$pending = $this->pendingData[$frame->sessionId] ?? '';
unset($this->pendingData[$frame->sessionId]);
if ($pending !== '') {
$this->send(new Frame(FrameType::DATA, $frame->sessionId, ['data_raw' => $pending]));
$this->sendData($frame->sessionId, $pending);
}
}
@@ -624,26 +630,49 @@ final class AgentClient
unset($this->initialBuffers[$connection->id]);
unset($this->connectionStages[$connection->id]);
unset($this->pausedClientsForPop[$connection->id], $this->clientsPausingPop[$connection->id]);
$suppressCloseFrame = isset($this->suppressClientCloseFrames[$connection->id]);
unset($this->suppressClientCloseFrames[$connection->id]);
if (!isset($this->connectionSessionIds[$connection->id])) {
return;
}
$sessionId = $this->connectionSessionIds[$connection->id];
unset($this->pendingClientCloses[$sessionId]);
unset($this->connectionSessionIds[$connection->id]);
unset($this->clients[$sessionId], $this->sessionStates[$sessionId], $this->pendingData[$sessionId], $this->sessionIngressProtocols[$sessionId]);
if ($suppressCloseFrame) {
return;
}
$this->send(new Frame(FrameType::CLOSE, $sessionId, ['reason' => 'client_closed']));
}
private function closeClient(?string $sessionId): void
private function closeClient(?string $sessionId, bool $graceful = false): void
{
if ($sessionId === null || !isset($this->clients[$sessionId])) {
return;
}
$connection = $this->clients[$sessionId];
if ($graceful && $connection->getSendBufferQueueSize() > 0) {
$this->pendingClientCloses[$sessionId] = true;
$this->suppressClientCloseFrames[$connection->id] = true;
return;
}
$this->finalizeClientClose($sessionId);
}
private function finalizeClientClose(string $sessionId): void
{
if (!isset($this->clients[$sessionId])) {
return;
}
$connection = $this->clients[$sessionId];
unset($this->clients[$sessionId], $this->sessionStates[$sessionId], $this->pendingData[$sessionId], $this->sessionIngressProtocols[$sessionId]);
unset($this->connectionSessionIds[$connection->id]);
unset($this->pausedClientsForPop[$connection->id], $this->clientsPausingPop[$connection->id]);
unset($this->pausedClientsForPop[$connection->id], $this->clientsPausingPop[$connection->id], $this->pendingClientCloses[$sessionId]);
$this->suppressClientCloseFrames[$connection->id] = true;
$connection->close();
}
@@ -801,6 +830,21 @@ final class AgentClient
return $this->pop?->send(FrameCodec::encode($frame));
}
private function sendData(string $sessionId, string $data): bool
{
$length = strlen($data);
for ($offset = 0; $offset < $length; $offset += self::DATA_CHUNK_SIZE) {
$sent = $this->send(new Frame(FrameType::DATA, $sessionId, [
'data_raw' => substr($data, $offset, self::DATA_CHUNK_SIZE),
]));
if ($sent === false) {
return false;
}
}
return true;
}
private function resumeClientsForPop(): void
{
foreach ($this->pausedClientsForPop as $connectionId => $client) {
@@ -822,4 +866,13 @@ final class AgentClient
$this->pop?->resumeRecv();
}
}
private function onClientBufferDrain(TcpConnection $connection): void
{
$this->resumePopForClient($connection);
$sessionId = $this->connectionSessionIds[$connection->id] ?? null;
if (is_string($sessionId) && isset($this->pendingClientCloses[$sessionId])) {
$this->finalizeClientClose($sessionId);
}
}
}
+21 -10
View File
@@ -26,6 +26,7 @@ final class AgentListener
{
private const MAX_SEND_BUFFER = 32 * 1024 * 1024;
private const BACKPRESSURE_HIGH_WATERMARK = 16 * 1024 * 1024;
private const DATA_CHUNK_SIZE = 256 * 1024;
/** @var array<int, FrameParser> */
private array $parsers = [];
@@ -96,18 +97,16 @@ final class AgentListener
}
$nodeId = $this->connectionNodeIds[$connection->id] ?? null;
if (!is_string($nodeId) || $this->nodes->get($nodeId) === null) {
$node = is_string($nodeId) ? $this->nodes->get($nodeId) : null;
if (!is_string($nodeId) || $node === null) {
$this->send($connection, new Frame(FrameType::ERROR, $frame->sessionId, ['reason' => 'invalid_auth']));
$connection->close();
return;
}
$node->lastHeartbeat = time();
if ($frame->type === FrameType::PING) {
$node = $this->nodes->get($nodeId);
if ($node !== null) {
$node->lastHeartbeat = time();
$node->activeSessions = (int)($frame->payload['active_sessions'] ?? $node->activeSessions);
}
$node->activeSessions = (int)($frame->payload['active_sessions'] ?? $node->activeSessions);
$this->send($connection, new Frame(FrameType::PONG, null, ['timestamp' => time()]));
return;
}
@@ -244,10 +243,7 @@ final class AgentListener
};
$target->onMessage = function (AsyncTcpConnection $target, string $data) use ($session, $agentConnection): void {
$session->bytesTargetToClient += strlen($data);
$sent = $this->send($agentConnection, new Frame(FrameType::DATA, $session->sessionId, [
'data_raw' => $data,
]));
if ($sent === false) {
if (!$this->sendData($agentConnection, $session->sessionId, $data)) {
$this->closeSession($session, 'failed', 'agent_buffer_overflow');
return;
}
@@ -396,6 +392,21 @@ final class AgentListener
return $connection->send(FrameCodec::encode($frame));
}
private function sendData(TcpConnection $connection, string $sessionId, string $data): bool
{
$length = strlen($data);
for ($offset = 0; $offset < $length; $offset += self::DATA_CHUNK_SIZE) {
$sent = $this->send($connection, new Frame(FrameType::DATA, $sessionId, [
'data_raw' => substr($data, $offset, self::DATA_CHUNK_SIZE),
]));
if ($sent === false) {
return false;
}
}
return true;
}
private function resumeTargetsForAgent(TcpConnection $connection): void
{
foreach ($this->pausedTargetsByAgent[$connection->id] ?? [] as $sessionId => $target) {