subprocess-util/src/subprocess_util/exec_send_chunks_passive.py
2023-03-19 16:13:54 +01:00

55 lines
2.0 KiB
Python

import socket
from pathlib import Path
from subprocess_util.common import Message1, Message2
from subprocess_util.exec_produce_chunks import execute_produce_chunks
from subprocess_util.unix_sock_input import accept_loop_until_message
def exec_send_chunks_passive(
command: list[str],
source_chunk_dir: Path,
chunk_size: int = 128 * 1024 * 1024,
) -> int:
def get_source_chunk_path(chunk_no: int) -> Path:
return source_chunk_dir.joinpath(f'{chunk_no}')
def get_source_created_path(chunk_no: int) -> Path:
return source_chunk_dir.joinpath(f'{chunk_no}.COMPLETE')
def get_source_transferred_path() -> Path:
return source_chunk_dir.joinpath('SOCKET')
def handle_chunk(chunk_no: int, last_chunk: bool):
source_chunk_path = get_source_chunk_path(chunk_no)
print(f'Handling chunk {source_chunk_path}')
# Create in_socket
chunk_transferred_sock_path = get_source_transferred_path()
chunk_transferred_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
chunk_transferred_sock.bind(str(chunk_transferred_sock_path))
chunk_transferred_sock.listen(1)
# Inform receiving side, that the next chunk can be transferred.
chunk_created_path = get_source_created_path(chunk_no)
chunk_created_path.write_text(Message1.EOF.value if last_chunk else Message1.OK.value)
# Read from in_socket
# -> Receiving side informs us after the chunk has been transferred.
_command = accept_loop_until_message(chunk_transferred_sock, [Message2.OK_BINARY.value])
# Delete local chunk after it has been transferred.
source_chunk_path.unlink(missing_ok=False)
# Close in_socket
chunk_transferred_sock.close()
chunk_transferred_sock_path.unlink(missing_ok=False)
returncode = execute_produce_chunks(
command=command,
get_chunk_path=get_source_chunk_path,
handle_chunk=handle_chunk,
chunk_size=chunk_size,
)
return returncode