mirror of
https://github.com/langgenius/dify.git
synced 2026-05-13 08:57:28 +08:00
feat(docker-demuxer): implement producer-consumer pattern for stream demultiplexing
- Introduced threading to handle Docker's stdout/stderr streams, improving thread safety and preventing race conditions. - Replaced buffer-based reading with queue-based reading for stdout and stderr. - Updated read methods to handle errors and end-of-stream conditions more gracefully. - Enhanced documentation to reflect changes in the demuxing process.
This commit is contained in:
parent
c2e5081437
commit
beefff3d48
@ -1,10 +1,13 @@
|
|||||||
|
import logging
|
||||||
import socket
|
import socket
|
||||||
import tarfile
|
import tarfile
|
||||||
|
import threading
|
||||||
from collections.abc import Mapping, Sequence
|
from collections.abc import Mapping, Sequence
|
||||||
from enum import IntEnum, StrEnum
|
from enum import IntEnum, StrEnum
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from pathlib import PurePosixPath
|
from pathlib import PurePosixPath
|
||||||
|
from queue import Queue
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
@ -40,7 +43,7 @@ class DockerStreamType(IntEnum):
|
|||||||
|
|
||||||
class DockerDemuxer:
|
class DockerDemuxer:
|
||||||
"""
|
"""
|
||||||
Demultiplexes Docker's combined stdout/stderr stream.
|
Demultiplexes Docker's combined stdout/stderr stream using producer-consumer pattern.
|
||||||
|
|
||||||
Docker exec with tty=False sends stdout and stderr over a single socket,
|
Docker exec with tty=False sends stdout and stderr over a single socket,
|
||||||
each frame prefixed with an 8-byte header:
|
each frame prefixed with an 8-byte header:
|
||||||
@ -48,58 +51,56 @@ class DockerDemuxer:
|
|||||||
- Bytes 1-3: reserved (zeros)
|
- Bytes 1-3: reserved (zeros)
|
||||||
- Bytes 4-7: payload size (big-endian uint32)
|
- Bytes 4-7: payload size (big-endian uint32)
|
||||||
|
|
||||||
This class reads frames and routes them to separate stdout/stderr buffers.
|
THREAD SAFETY:
|
||||||
Without demuxing, output contains binary garbage like:
|
A single background thread reads frames from the socket and dispatches payloads
|
||||||
b'\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x0eHello World\\n'
|
to thread-safe queues. This avoids race conditions where multiple threads
|
||||||
|
calling _read_next_frame() simultaneously caused frame header/body corruption,
|
||||||
|
resulting in incomplete stdout/stderr output.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
_HEADER_SIZE = 8
|
_HEADER_SIZE = 8
|
||||||
|
|
||||||
def __init__(self, sock: socket.SocketIO):
|
def __init__(self, sock: socket.SocketIO):
|
||||||
self._sock = sock
|
self._sock = sock
|
||||||
self._stdout_buf = bytearray()
|
self._stdout_queue: Queue[bytes | None] = Queue()
|
||||||
self._stderr_buf = bytearray()
|
self._stderr_queue: Queue[bytes | None] = Queue()
|
||||||
self._eof = False
|
|
||||||
self._closed = False
|
self._closed = False
|
||||||
|
self._error: BaseException | None = None
|
||||||
|
|
||||||
def read_stdout(self, n: int) -> bytes:
|
self._demux_thread = threading.Thread(
|
||||||
return self._read_from_buffer(self._stdout_buf, DockerStreamType.STDOUT, n)
|
target=self._demux_loop,
|
||||||
|
daemon=True,
|
||||||
|
name="docker-demuxer",
|
||||||
|
)
|
||||||
|
self._demux_thread.start()
|
||||||
|
|
||||||
def read_stderr(self, n: int) -> bytes:
|
def _demux_loop(self) -> None:
|
||||||
return self._read_from_buffer(self._stderr_buf, DockerStreamType.STDERR, n)
|
try:
|
||||||
|
while not self._closed:
|
||||||
|
header = self._read_exact(self._HEADER_SIZE)
|
||||||
|
if not header or len(header) < self._HEADER_SIZE:
|
||||||
|
break
|
||||||
|
|
||||||
def _read_from_buffer(self, buffer: bytearray, target_type: DockerStreamType, n: int) -> bytes:
|
frame_type = header[0]
|
||||||
while len(buffer) < n and not self._eof:
|
payload_size = int.from_bytes(header[4:8], "big")
|
||||||
self._read_next_frame()
|
|
||||||
|
|
||||||
if not buffer:
|
if payload_size == 0:
|
||||||
raise TransportEOFError("End of demuxed stream")
|
continue
|
||||||
|
|
||||||
result = bytes(buffer[:n])
|
payload = self._read_exact(payload_size)
|
||||||
del buffer[:n]
|
if not payload:
|
||||||
return result
|
break
|
||||||
|
|
||||||
def _read_next_frame(self) -> None:
|
if frame_type == DockerStreamType.STDOUT:
|
||||||
header = self._read_exact(self._HEADER_SIZE)
|
self._stdout_queue.put(payload)
|
||||||
if not header or len(header) < self._HEADER_SIZE:
|
elif frame_type == DockerStreamType.STDERR:
|
||||||
self._eof = True
|
self._stderr_queue.put(payload)
|
||||||
return
|
|
||||||
|
|
||||||
frame_type = header[0]
|
except BaseException as e:
|
||||||
payload_size = int.from_bytes(header[4:8], "big")
|
self._error = e
|
||||||
|
finally:
|
||||||
if payload_size == 0:
|
self._stdout_queue.put(None)
|
||||||
return
|
self._stderr_queue.put(None)
|
||||||
|
|
||||||
payload = self._read_exact(payload_size)
|
|
||||||
if not payload:
|
|
||||||
self._eof = True
|
|
||||||
return
|
|
||||||
|
|
||||||
if frame_type == DockerStreamType.STDOUT:
|
|
||||||
self._stdout_buf.extend(payload)
|
|
||||||
elif frame_type == DockerStreamType.STDERR:
|
|
||||||
self._stderr_buf.extend(payload)
|
|
||||||
|
|
||||||
def _read_exact(self, size: int) -> bytes:
|
def _read_exact(self, size: int) -> bytes:
|
||||||
data = bytearray()
|
data = bytearray()
|
||||||
@ -115,18 +116,50 @@ class DockerDemuxer:
|
|||||||
return bytes(data) if data else b""
|
return bytes(data) if data else b""
|
||||||
return bytes(data)
|
return bytes(data)
|
||||||
|
|
||||||
|
def read_stdout(self) -> bytes:
|
||||||
|
return self._read_from_queue(self._stdout_queue)
|
||||||
|
|
||||||
|
def read_stderr(self) -> bytes:
|
||||||
|
return self._read_from_queue(self._stderr_queue)
|
||||||
|
|
||||||
|
def _read_from_queue(self, queue: Queue[bytes | None]) -> bytes:
|
||||||
|
if self._error:
|
||||||
|
raise TransportEOFError(f"Demuxer error: {self._error}") from self._error
|
||||||
|
|
||||||
|
chunk = queue.get()
|
||||||
|
if chunk is None:
|
||||||
|
if self._error:
|
||||||
|
raise TransportEOFError(f"Demuxer error: {str(self._error)}")
|
||||||
|
raise TransportEOFError("End of demuxed stream")
|
||||||
|
return chunk
|
||||||
|
|
||||||
def close(self) -> None:
|
def close(self) -> None:
|
||||||
if not self._closed:
|
if not self._closed:
|
||||||
self._closed = True
|
self._closed = True
|
||||||
self._sock.close()
|
try:
|
||||||
|
self._sock.close()
|
||||||
|
except Exception:
|
||||||
|
logging.error("Failed to close Docker demuxer socket", exc_info=True)
|
||||||
|
|
||||||
|
|
||||||
class DemuxedStdoutReader(TransportReadCloser):
|
class DemuxedStdoutReader(TransportReadCloser):
|
||||||
def __init__(self, demuxer: DockerDemuxer):
|
def __init__(self, demuxer: DockerDemuxer):
|
||||||
self._demuxer = demuxer
|
self._demuxer = demuxer
|
||||||
|
self._buffer = bytearray()
|
||||||
|
|
||||||
def read(self, n: int) -> bytes:
|
def read(self, n: int) -> bytes:
|
||||||
return self._demuxer.read_stdout(n)
|
if self._buffer:
|
||||||
|
data = bytes(self._buffer[:n])
|
||||||
|
del self._buffer[:n]
|
||||||
|
if data:
|
||||||
|
return data
|
||||||
|
|
||||||
|
chunk = self._demuxer.read_stdout()
|
||||||
|
if len(chunk) <= n:
|
||||||
|
return chunk
|
||||||
|
|
||||||
|
self._buffer.extend(chunk[n:])
|
||||||
|
return chunk[:n]
|
||||||
|
|
||||||
def close(self) -> None:
|
def close(self) -> None:
|
||||||
self._demuxer.close()
|
self._demuxer.close()
|
||||||
@ -135,9 +168,21 @@ class DemuxedStdoutReader(TransportReadCloser):
|
|||||||
class DemuxedStderrReader(TransportReadCloser):
|
class DemuxedStderrReader(TransportReadCloser):
|
||||||
def __init__(self, demuxer: DockerDemuxer):
|
def __init__(self, demuxer: DockerDemuxer):
|
||||||
self._demuxer = demuxer
|
self._demuxer = demuxer
|
||||||
|
self._buffer = bytearray()
|
||||||
|
|
||||||
def read(self, n: int) -> bytes:
|
def read(self, n: int) -> bytes:
|
||||||
return self._demuxer.read_stderr(n)
|
if self._buffer:
|
||||||
|
data = bytes(self._buffer[:n])
|
||||||
|
del self._buffer[:n]
|
||||||
|
if data:
|
||||||
|
return data
|
||||||
|
|
||||||
|
chunk = self._demuxer.read_stderr()
|
||||||
|
if len(chunk) <= n:
|
||||||
|
return chunk
|
||||||
|
|
||||||
|
self._buffer.extend(chunk[n:])
|
||||||
|
return chunk[:n]
|
||||||
|
|
||||||
def close(self) -> None:
|
def close(self) -> None:
|
||||||
self._demuxer.close()
|
self._demuxer.close()
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user