class IPCHandler(logging.Handler):
logger = get_logger('trinity._utils.logging.IPCHandler')
def __init__(self, sock: socket.socket):
self._socket = BufferedSocket(sock)
super().__init__()
def __enter__(self) -> None:
pass
def __exit__(self,
exc_type: Type[BaseException],
exc_value: BaseException,
exc_tb: TracebackType) -> None:
self._socket.close()
@classmethod
def connect(cls: Type[THandler], path: Path) -> THandler:
wait_for_ipc(path)
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
cls.logger.debug("Opened connection to %s: %s", path, s)
s.connect(str(path))
return cls(s)
def prepare(self, record: logging.LogRecord) -> logging.LogRecord:
msg = self.format(record)
new_record = copy.copy(record)
new_record.message = msg
new_record.msg = msg
new_record.args = None
new_record.exc_info = None
new_record.exc_text = None
return new_record
def emit(self, record: logging.LogRecord) -> None:
try:
msg_data = pickle.dumps(self.prepare(record))
msg_length_data = len(msg_data).to_bytes(4, 'big')
self._socket.sendall(msg_length_data + msg_data)
except Exception:
self.handleError(record)
class DBClient(BaseAtomicDB):
logger = logging.getLogger('trinity.db.client.DBClient')
def __init__(self, sock: socket.socket):
self._socket = BufferedSocket(sock)
self._lock = threading.Lock()
def __enter__(self) -> None:
self._socket.__enter__()
def __exit__(self, exc_type: Type[BaseException], exc_value: BaseException,
exc_tb: TracebackType) -> None:
self._socket.__exit__(exc_type, exc_value, exc_tb)
def __getitem__(self, key: bytes) -> bytes:
with self._lock:
self._socket.sendall(GET.value +
len(key).to_bytes(LEN_BYTES, 'little') + key)
result_byte = self._socket.read_exactly(1)
if result_byte == SUCCESS_BYTE:
value_size_data = self._socket.read_exactly(LEN_BYTES)
value = self._socket.read_exactly(
int.from_bytes(value_size_data, 'little'))
return value
elif result_byte == FAIL_BYTE:
raise KeyError(key)
else:
raise Exception(f"Unknown result byte: {result_byte.hex}")
def __setitem__(self, key: bytes, value: bytes) -> None:
with self._lock:
self._socket.sendall(SET.value +
struct.pack('<II', len(key), len(value)) +
key + value)
Result(self._socket.read_exactly(1))
def __delitem__(self, key: bytes) -> None:
with self._lock:
self._socket.sendall(DELETE.value +
len(key).to_bytes(4, 'little') + key)
result_byte = self._socket.read_exactly(1)
if result_byte == SUCCESS_BYTE:
return
elif result_byte == FAIL_BYTE:
raise KeyError(key)
else:
raise Exception(f"Unknown result byte: {result_byte.hex}")
def _exists(self, key: bytes) -> bool:
with self._lock:
self._socket.sendall(EXISTS.value +
len(key).to_bytes(4, 'little') + key)
result_byte = self._socket.read_exactly(1)
if result_byte == SUCCESS_BYTE:
return True
elif result_byte == FAIL_BYTE:
return False
else:
raise Exception(f"Unknown result byte: {result_byte.hex}")
@contextlib.contextmanager
def atomic_batch(self) -> Iterator['AtomicBatch']:
batch = AtomicBatch(self)
yield batch
diff = batch.finalize()
pending_deletes = diff.deleted_keys()
pending_kv_pairs = diff.pending_items()
kv_pair_count = len(pending_kv_pairs)
delete_count = len(pending_deletes)
kv_sizes = tuple(
len(item) for item in itertools.chain(*pending_kv_pairs))
delete_sizes = tuple(len(key) for key in pending_deletes)
# We encode all of the *sizes* in one shot using `struct.pack` and this
# dynamically constructed format string.
fmt_str = '<II' + 'I' * (len(kv_sizes) + len(pending_deletes))
kv_pair_count_and_size_data = struct.pack(
fmt_str,
kv_pair_count,
delete_count,
*kv_sizes,
*delete_sizes,
)
kv_and_delete_data = b''.join(
itertools.chain(*pending_kv_pairs, pending_deletes))
with self._lock:
self._socket.sendall(ATOMIC_BATCH.value +
kv_pair_count_and_size_data +
kv_and_delete_data)
Result(self._socket.read_exactly(1))
def close(self) -> None:
try:
self._socket.shutdown(socket.SHUT_WR)
except OSError as e:
# on mac OS this can result in the following error:
# OSError: [Errno 57] Socket is not connected
if e.errno != errno.ENOTCONN:
raise
self._socket.close()
@classmethod
def connect(cls, path: pathlib.Path, timeout: int = 5) -> "DBClient":
wait_for_ipc(path, timeout)
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
cls.logger.debug("Opened connection to %s: %s", path, s)
s.connect(str(path))
return cls(s)