From e9f28e09ce9877b9b86f1c5714b4da73c25c10aa Mon Sep 17 00:00:00 2001 From: Gewu <89496957+RkGrit@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:30:45 +0800 Subject: [PATCH 1/3] fix(python): release dataset index file after mmap --- python/tests/test_dataset_index.py | 228 +++++++++++++++++++++++++++++ python/tsfile/dataset/index.py | 44 +++--- 2 files changed, 255 insertions(+), 17 deletions(-) diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index 09a52e3c0..4e5a4a344 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -716,3 +716,231 @@ def close_lease(): close_thread.join(timeout=2) assert query_done.is_set() assert close_done.is_set() + + +def _write_tiny_mapped_index(path): + source = path.with_suffix(".tsfile") + source.write_bytes(b"T" * 4096) + write_index_atomic( + str(path), build_sections_from_dataframe(_synthetic_dataframe(str(source))) + ) + + +def _target_inode(path): + stat = os.stat(path) + return stat.st_dev, stat.st_ino + + +def _target_inode_fd_count(path): + if not os.path.isdir("/proc/self/fd"): + pytest.skip("Linux /proc FD attribution is unavailable") + target = _target_inode(path) + count = 0 + for fd in os.listdir("/proc/self/fd"): + try: + stat = os.stat(f"/proc/self/fd/{fd}") + except FileNotFoundError: + continue + if (stat.st_dev, stat.st_ino) == target: + count += 1 + return count + + +def _target_inode_vma_count(path): + if not os.path.isfile("/proc/self/maps"): + pytest.skip("Linux /proc VMA attribution is unavailable") + target_device, target_inode = _target_inode(path) + target_device = f"{os.major(target_device):02x}:{os.minor(target_device):02x}" + target_inode = str(target_inode) + count = 0 + with open("/proc/self/maps", encoding="utf-8") as maps: + for line in maps: + fields = line.split(maxsplit=5) + if ( + len(fields) >= 5 + and fields[3] == target_device + and fields[4] == target_inode + ): + count += 1 + return count + + +def _assert_index_lookup(index): + table_id = index.find_table_ids("root")[0] + device_id = index.find_device_id(table_id, "root.") + column_id = index.find_column_id(table_id, "s1") + assert index.find_series_id(device_id, column_id) == 0 + + +def test_mapped_index_releases_original_fd_after_readonly_mmap(tmp_path): + output = tmp_path / "single.tsidx" + _write_tiny_mapped_index(output) + baseline_fds = _target_inode_fd_count(output) + baseline_vmas = _target_inode_vma_count(output) + + index = MappedDatasetIndex(str(output)) + try: + _assert_index_lookup(index) + assert _target_inode_fd_count(output) == baseline_fds + 1 + assert index._file is None + assert _target_inode_vma_count(output) == baseline_vmas + 1 + finally: + index.close() + + assert _target_inode_fd_count(output) == baseline_fds + assert _target_inode_vma_count(output) == baseline_vmas + + +def test_mapped_indices_keep_one_fd_and_live_vma_per_mapping(tmp_path): + paths = [tmp_path / f"index-{number}.tsidx" for number in range(16)] + for path in paths: + _write_tiny_mapped_index(path) + + baseline_vmas = {path: _target_inode_vma_count(path) for path in paths} + indices = [MappedDatasetIndex(str(path)) for path in paths] + try: + for path, index in zip(paths, indices): + _assert_index_lookup(index) + assert _target_inode_fd_count(path) == 1 + assert index._file is None + assert _target_inode_vma_count(path) == baseline_vmas[path] + 1 + finally: + for index in indices: + index.close() + for index in indices: + index.close() + + for path in paths: + assert _target_inode_fd_count(path) == 0 + assert _target_inode_vma_count(path) == baseline_vmas[path] + + +def test_mapped_index_mmap_failure_releases_target_resources(tmp_path, monkeypatch): + output = tmp_path / "mmap-failure.tsidx" + _write_tiny_mapped_index(output) + baseline_fds = _target_inode_fd_count(output) + baseline_vmas = _target_inode_vma_count(output) + + def fail_mmap(*_args, **_kwargs): + raise OSError("injected mmap failure") + + monkeypatch.setattr(index_module.mmap, "mmap", fail_mmap) + with pytest.raises(OSError, match="injected mmap failure"): + MappedDatasetIndex(str(output)) + + assert _target_inode_fd_count(output) == baseline_fds + assert _target_inode_vma_count(output) == baseline_vmas + + +@pytest.mark.parametrize( + ("trust_index", "method_name"), + [(False, "_validate"), (True, "_map_entries_without_validation")], +) +def test_mapped_index_post_mmap_failure_releases_target_resources( + tmp_path, monkeypatch, trust_index, method_name +): + output = tmp_path / f"post-mmap-{trust_index}.tsidx" + _write_tiny_mapped_index(output) + baseline_fds = _target_inode_fd_count(output) + baseline_vmas = _target_inode_vma_count(output) + + def fail_after_mmap(*_args, **_kwargs): + raise RuntimeError("injected post-mmap failure") + + monkeypatch.setattr(MappedDatasetIndex, method_name, fail_after_mmap) + with pytest.raises(RuntimeError, match="injected post-mmap failure"): + MappedDatasetIndex(str(output), trust_index=trust_index) + + assert _target_inode_fd_count(output) == baseline_fds + assert _target_inode_vma_count(output) == baseline_vmas + + +def test_mapped_index_close_is_idempotent_and_closes_later_resources_after_error(): + events = [] + + class _View: + def release(self): + events.append("view.release") + raise RuntimeError("view release failed") + + class _Mmap: + def close(self): + events.append("mmap.close") + + class _File: + def close(self): + events.append("file.close") + + index = MappedDatasetIndex.__new__(MappedDatasetIndex) + index._view = _View() + index._mmap = _Mmap() + index._file = _File() + + with pytest.raises(RuntimeError, match="view release failed"): + index.close() + assert events == ["view.release", "mmap.close", "file.close"] + assert index._view is None + assert index._mmap is None + assert index._file is None + index.close() + + +def _linux_proc_resource_attribution_available(): + return ( + hasattr(os, "fork") + and os.path.isdir("/proc/self/fd") + and os.path.isfile("/proc/self/maps") + ) + + +def test_linux_proc_resource_attribution_guard_requires_proc(monkeypatch): + monkeypatch.setattr(os, "fork", lambda: None, raising=False) + monkeypatch.setattr(os.path, "isdir", lambda path: path != "/proc/self/fd") + monkeypatch.setattr(os.path, "isfile", lambda path: True) + assert not _linux_proc_resource_attribution_available() + + +def test_mapped_index_inherited_mapping_survives_child_double_close(tmp_path): + if not _linux_proc_resource_attribution_available(): + pytest.skip("Linux fork /proc resource attribution is unavailable") + output = tmp_path / "fork.tsidx" + _write_tiny_mapped_index(output) + index = MappedDatasetIndex(str(output)) + read_fd, write_fd = os.pipe() + child_pid = os.fork() + if child_pid == 0: + os.close(read_fd) + try: + _assert_index_lookup(index) + assert index._file is None + assert _target_inode_fd_count(output) == 1 + index.close() + index.close() + assert _target_inode_fd_count(output) == 0 + assert _target_inode_vma_count(output) == 0 + os.write(write_fd, b"OK") + status = 0 + except BaseException as exc: + os.write(write_fd, f"{type(exc).__name__}: {exc}".encode()) + status = 1 + finally: + os.close(write_fd) + os._exit(status) + + os.close(write_fd) + child_message = os.read(read_fd, 4096) + os.close(read_fd) + waited_pid, child_status = os.waitpid(child_pid, 0) + try: + assert waited_pid == child_pid + assert os.WIFEXITED(child_status), child_message.decode() + assert os.WEXITSTATUS(child_status) == 0, child_message.decode() + assert child_message == b"OK" + _assert_index_lookup(index) + assert _target_inode_fd_count(output) == 1 + finally: + index.close() + index.close() + + assert _target_inode_fd_count(output) == 0 + assert _target_inode_vma_count(output) == 0 diff --git a/python/tsfile/dataset/index.py b/python/tsfile/dataset/index.py index 790fcd077..6ed7cd859 100644 --- a/python/tsfile/dataset/index.py +++ b/python/tsfile/dataset/index.py @@ -242,8 +242,11 @@ def __init__( ): self.path = path self._trusted = bool(trust_index) - self._file = open(path, "rb") + self._file = None + self._mmap = None + self._view = None try: + self._file = open(path, "rb") stat = os.fstat(self._file.fileno()) self.identity = ( stat.st_dev, @@ -252,6 +255,9 @@ def __init__( stat.st_mtime_ns, ) self._mmap = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ) + if os.name == "posix" and os.uname().sysname == "Linux": + self._file.close() + self._file = None self._view = memoryview(self._mmap) self._entries = ( self._map_entries_without_validation() @@ -259,15 +265,27 @@ def __init__( else self._validate(verify_sections) ) except Exception: - if getattr(self, "_view", None) is not None: - self._view.release() - self._view = None - if getattr(self, "_mmap", None) is not None: - self._mmap.close() - self._mmap = None - self._file.close() + self._close_owned_resources() raise + def _close_owned_resources(self): + view = self._view + self._view = None + try: + if view is not None: + view.release() + finally: + mapped = self._mmap + self._mmap = None + try: + if mapped is not None: + mapped.close() + finally: + file_object = self._file + self._file = None + if file_object is not None: + file_object.close() + def _map_entries_without_validation(self): """Read v1 section descriptors without checking their contents. @@ -374,15 +392,7 @@ def _validate(self, verify_sections: bool): return entries def close(self): - if getattr(self, "_view", None) is not None: - self._view.release() - self._view = None - if getattr(self, "_mmap", None) is not None: - self._mmap.close() - self._mmap = None - if getattr(self, "_file", None) is not None: - self._file.close() - self._file = None + self._close_owned_resources() def __enter__(self): return self From 3d21b820140e86ef3bf25e2ce7fa62c99c38068f Mon Sep 17 00:00:00 2001 From: Gewu <89496957+RkGrit@users.noreply.github.com> Date: Tue, 8 Sep 2026 14:49:28 +0800 Subject: [PATCH 2/3] fix(python): bound dataset runtime resource ownership --- python/tests/test_dataset_index.py | 826 ++++++++++++++++++++++++++++- python/tsfile/dataset/dataframe.py | 46 +- python/tsfile/dataset/runtime.py | 487 ++++++++++++++--- 3 files changed, 1273 insertions(+), 86 deletions(-) diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index 4e5a4a344..9b3eab077 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -622,18 +622,18 @@ def test_prepared_query_reads_nullable_offset_window_in_arrow_batches(tmp_path): series = runtime.index.record(LOGICAL_SERIES, 0) span = runtime.index.record(SERIES_FILE_SPAN, series[2]) with runtime.readers.acquire(0) as reader: - prepared = runtime.prepared.get(0, span[2], reader) - with reader.query_prepared(prepared, offset=1, limit=7) as result: - batches = [] - while True: - batch = result.read_arrow_batch() - if batch is None: - break - batches.append(batch) - with reader.query_prepared( - prepared, start_time=100, end_time=200 - ) as empty_result: - assert empty_result.read_arrow_batch() is None + with runtime.prepared.acquire(0, span[2], reader) as prepared: + with reader.query_prepared(prepared, offset=1, limit=7) as result: + batches = [] + while True: + batch = result.read_arrow_batch() + if batch is None: + break + batches.append(batch) + with reader.query_prepared( + prepared, start_time=100, end_time=200 + ) as empty_result: + assert empty_result.read_arrow_batch() is None assert batches table = pa.concat_tables(batches) @@ -944,3 +944,805 @@ def test_mapped_index_inherited_mapping_survives_child_double_close(tmp_path): assert _target_inode_fd_count(output) == 0 assert _target_inode_vma_count(output) == 0 + + +class _LifecycleAbort(BaseException): + pass + + +class _PreparedIndexStub: + @staticmethod + def record(section_type, record_id): + if section_type == SERIES_LOCATOR: + return (0, 0, record_id * 10, 8, 1) + if section_type == DEVICE_FILE_SPAN: + return (0, 0, 100, 16, 1, 0, 2) + if section_type == TSFILE_RECORD: + return (0, 0, 4096, 12345) + raise AssertionError(section_type) + + +class _PreparedHandleStub: + def __init__( + self, locator_id, close_error=None, time_owner=None, close_events=None + ): + self.locator_id = locator_id + self.close_error = close_error + self.time_owner = time_owner + self.close_events = close_events + self.close_calls = 0 + + def close(self): + self.close_calls += 1 + if self.close_events is not None: + self.close_events.append(f"prepared.close:{self.locator_id}") + if self.close_error is not None: + raise self.close_error + + +class _PreparedReaderStub: + def __init__( + self, + *, + abort_once=False, + abort_query=False, + close_error_ids=(), + close_events=None, + query_result=None, + ): + self.abort_once = abort_once + self.abort_query = abort_query + self.close_error_ids = set(close_error_ids) + self.close_events = close_events + self.query_result = query_result + self.prepared = [] + + def prepare_series(self, locator, time_owner=None): + if self.abort_once: + self.abort_once = False + raise _LifecycleAbort("prepare interrupted") + locator_id = locator[4] + prepared = _PreparedHandleStub( + locator_id, + close_error=( + RuntimeError(f"close {locator_id}") + if locator_id in self.close_error_ids + else None + ), + time_owner=time_owner, + close_events=self.close_events, + ) + self.prepared.append(prepared) + return prepared + + def query_prepared(self, _prepared, **_kwargs): + if self.abort_query: + raise _LifecycleAbort("query interrupted") + return self.query_result + + +def test_prepared_cache_evicts_lru_and_plateaus_at_configured_cap(): + reader = _PreparedReaderStub() + cache = runtime_module.PreparedSeriesCache(_PreparedIndexStub(), max_entries=2) + + with cache.acquire(0, 0, reader) as first: + assert first.locator_id == 0 + with cache.acquire(0, 1, reader): + pass + with cache.acquire(0, 0, reader) as reused: + assert reused is first + with cache.acquire(0, 2, reader): + pass + + assert cache.size == 2 + assert [entry.close_calls for entry in reader.prepared] == [0, 1, 0] + cache.close() + assert [entry.close_calls for entry in reader.prepared] == [1, 1, 1] + + +def test_prepared_cache_default_is_finite_and_plateaus_without_environment( + monkeypatch, +): + monkeypatch.delenv("TSFILE_DATAFRAME_MAX_PREPARED_SERIES", raising=False) + reader = _PreparedReaderStub() + cache = runtime_module.PreparedSeriesCache(_PreparedIndexStub()) + + for locator_id in range(4097): + with cache.acquire(0, locator_id, reader): + pass + + assert cache.max_entries == 4096 + assert cache.size == 4096 + assert reader.prepared[0].close_calls == 1 + assert all(entry.close_calls == 0 for entry in reader.prepared[1:]) + cache.close() + assert all(entry.close_calls == 1 for entry in reader.prepared) + + +def test_prepared_cache_keeps_active_lru_until_lease_release(): + reader = _PreparedReaderStub() + cache = runtime_module.PreparedSeriesCache(_PreparedIndexStub(), max_entries=1) + + with cache.acquire(0, 0, reader) as active: + with cache.acquire(0, 1, reader): + assert cache.size == 2 + assert active.close_calls == 0 + assert cache.size == 1 + assert active.close_calls == 0 + + with cache.acquire(0, 2, reader): + pass + assert active.close_calls == 1 + assert cache.size == 1 + cache.close() + + +@pytest.mark.parametrize("max_entries", [0, 1, 2]) +def test_prepared_cache_owner_graph_plateaus_without_dependency_chains(max_entries): + reader = _PreparedReaderStub() + cache = runtime_module.PreparedSeriesCache( + _PreparedIndexStub(), max_entries=max_entries + ) + + for locator_id in range(1, 33): + with cache.acquire(0, locator_id - 1, reader) as owner: + with cache.acquire(0, locator_id, reader, time_owner=owner): + pass + + reachable = {} + pending = [entry.prepared for entry in cache._entries.values()] + while pending: + prepared = pending.pop() + if id(prepared) in reachable: + continue + reachable[id(prepared)] = prepared + if prepared.time_owner is not None: + pending.append(prepared.time_owner) + + assert cache.size <= cache.max_entries + assert len(reachable) <= cache.max_entries + assert all( + prepared.time_owner is None or prepared.time_owner.time_owner is None + for prepared in reachable.values() + ) + + cache.close() + assert all(prepared.close_calls == 1 for prepared in reader.prepared) + + +def test_prepared_cache_zero_retention_closes_each_handle_exactly_once(): + reader = _PreparedReaderStub() + cache = runtime_module.PreparedSeriesCache(_PreparedIndexStub(), max_entries=0) + + for _ in range(3): + with cache.acquire(0, 0, reader): + assert cache.size == 1 + assert cache.size == 0 + + assert len(reader.prepared) == 3 + assert [entry.close_calls for entry in reader.prepared] == [1, 1, 1] + cache.close() + assert [entry.close_calls for entry in reader.prepared] == [1, 1, 1] + + +def test_prepared_cache_recovers_single_flight_after_prepare_baseexception(): + reader = _PreparedReaderStub(abort_once=True) + cache = runtime_module.PreparedSeriesCache(_PreparedIndexStub(), max_entries=0) + + with pytest.raises(_LifecycleAbort, match="prepare interrupted"): + with cache.acquire(0, 0, reader): + pass + + assert not cache._loading + with cache.acquire(0, 0, reader) as prepared: + assert prepared.locator_id == 0 + assert prepared.close_calls == 1 + cache.close() + + +def test_prepared_cache_close_is_idempotent_and_closes_past_errors(): + reader = _PreparedReaderStub(close_error_ids={0}) + cache = runtime_module.PreparedSeriesCache(_PreparedIndexStub()) + with cache.acquire(0, 0, reader) as first: + pass + with cache.acquire(0, 1, reader) as second: + pass + + with pytest.raises(RuntimeError, match="close 0"): + cache.close() + assert first.close_calls == 1 + assert second.close_calls == 1 + + cache.close() + assert first.close_calls == 1 + assert second.close_calls == 1 + + +def test_runtime_discard_after_fork_is_lock_free_exact_once_and_resilient(): + from concurrent.futures import thread as thread_pool_module + import queue + + close_calls = [] + + class _PoisonCondition: + def __enter__(self): + raise AssertionError("inherited condition must not be acquired") + + class _Discardable: + def __init__(self, name, fail=False): + self.name = name + self.fail = fail + + def discard_after_fork(self): + close_calls.append(self.name) + if self.fail: + raise RuntimeError(f"discard {self.name}") + + class _Index: + def close(self): + close_calls.append("index") + + class _Executor: + def __init__(self): + self._shutdown = False + self._work_queue = queue.SimpleQueue() + self._threads = {_InheritedThread()} + + def shutdown(self, **_kwargs): + raise AssertionError("an inherited executor must not be shut down") + + class _InheritedThread: + def __init__(self): + self._target = object() + self._args = (object(),) + self._kwargs = {"inherited": True} + + @staticmethod + def join(): + pass + + runtime = object.__new__(runtime_module.DatasetRuntime) + runtime.creator_pid = os.getpid() + 1 + runtime._condition = _PoisonCondition() + executor = _Executor() + inherited_thread = next(iter(executor._threads)) + thread_pool_module._threads_queues[inherited_thread] = executor._work_queue + runtime._query_executor = executor + runtime.prepared = _Discardable("prepared", fail=True) + runtime.readers = _Discardable("readers") + runtime.index = _Index() + catalog = SimpleNamespace( + runtime=runtime, index=runtime.index, _readers={0: object()} + ) + runtime.catalog = catalog + runtime._resources_closed = False + runtime._discarded_pid = None + runtime._accepting = True + runtime._torn_down = False + + with pytest.raises(RuntimeError, match="discard prepared"): + runtime.discard_after_fork() + + assert close_calls == ["prepared", "readers", "index"] + assert runtime._resources_closed + assert runtime._discarded_pid == os.getpid() + assert runtime._query_executor is None + assert executor._shutdown + assert executor._work_queue is None + assert executor._threads == set() + assert inherited_thread not in thread_pool_module._threads_queues + assert inherited_thread._target is None + assert inherited_thread._args == () + assert inherited_thread._kwargs == {} + assert runtime.prepared is None + assert runtime.readers is None + assert runtime.index is None + assert runtime.catalog is None + assert catalog.runtime is None + assert catalog.index is None + assert catalog._readers == {} + + runtime.discard_after_fork() + assert close_calls == ["prepared", "readers", "index"] + + +def test_reader_pool_never_evicts_an_active_reader(monkeypatch): + pool = runtime_module.ReaderSessionPool( + SimpleNamespace(), max_open_files=1, validate_generation=False + ) + sessions = {} + + class _Session: + def __init__(self, file_id): + self.reader = file_id + self.active_uses = 0 + self.close_calls = 0 + + def close(self): + self.close_calls += 1 + + def new_session(file_id): + session = _Session(file_id) + sessions[file_id] = session + return session + + monkeypatch.setattr(pool, "_new_session", new_session) + acquired_second = threading.Event() + + def acquire_second(): + with pool.acquire(1): + acquired_second.set() + + with pool.acquire(0): + thread = threading.Thread(target=acquire_second) + thread.start() + assert not acquired_second.wait(timeout=0.05) + assert sessions[0].close_calls == 0 + + thread.join(timeout=2) + assert acquired_second.is_set() + assert sessions[0].close_calls == 1 + pool.close() + assert sessions[1].close_calls == 1 + + +def test_runtime_partial_construction_rolls_back_all_resources(monkeypatch): + events = [] + + class _Index: + def __init__(self, *_args, **_kwargs): + events.append("index.open") + + def close(self): + events.append("index.close") + + class _Executor: + def __init__(self, *_args, **_kwargs): + events.append("executor.open") + + def shutdown(self, **kwargs): + events.append(("executor.shutdown", kwargs)) + + class _Readers: + def __init__(self, *_args, **_kwargs): + events.append("readers.open") + + def close(self): + events.append("readers.close") + + class _Prepared: + def __init__(self, *_args, **_kwargs): + events.append("prepared.open") + + def close(self): + events.append("prepared.close") + raise RuntimeError("prepared cleanup failed") + + class _Catalog: + def __init__(self, _runtime): + raise _LifecycleAbort("catalog construction interrupted") + + monkeypatch.setattr(runtime_module, "MappedDatasetIndex", _Index) + monkeypatch.setattr(runtime_module, "ThreadPoolExecutor", _Executor) + monkeypatch.setattr(runtime_module, "ReaderSessionPool", _Readers) + monkeypatch.setattr(runtime_module, "PreparedSeriesCache", _Prepared) + monkeypatch.setattr(runtime_module, "MappedDataFrameCatalog", _Catalog) + + with pytest.raises(_LifecycleAbort, match="catalog construction interrupted"): + runtime_module.DatasetRuntime("unused.tsidx", query_workers=2) + + assert events == [ + "index.open", + "executor.open", + "readers.open", + "prepared.open", + ("executor.shutdown", {"wait": True, "cancel_futures": True}), + "prepared.close", + "readers.close", + "index.close", + ] + + +@pytest.mark.parametrize("cache_size", [0, 2]) +def test_prepared_cache_follows_environment_cap(tmp_path, monkeypatch, cache_size): + source = tmp_path / "devices.tsfile" + _write_runtime_devices_file(source) + monkeypatch.setenv("TSFILE_DATAFRAME_MAX_PREPARED_SERIES", str(cache_size)) + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + for index in range(3): + series = dataframe[index] + np.testing.assert_array_equal( + series[:], np.array([index * 10.0, index * 10.0 + 1.0]) + ) + series.close() + + assert dataframe._runtime.prepared.max_entries == cache_size + assert dataframe._runtime.prepared.size == cache_size + + +class _ReaderPoolStub: + def __init__(self, reader, events): + self.reader = reader + self.events = events + self.acquire_calls = 0 + self.release_calls = 0 + + def acquire(self, _file_id): + pool = self + + class _Lease: + def __enter__(self): + pool.acquire_calls += 1 + return pool.reader + + def __exit__(self, *_args): + pool.release_calls += 1 + pool.events.append("reader.release") + + return _Lease() + + +class _FailingResultStub: + def __init__(self, events): + self.events = events + self.close_calls = 0 + + def __enter__(self): + return self + + def __exit__(self, *_args): + self.close_calls += 1 + self.events.append("result.close") + + @staticmethod + def read_arrow_batch(): + raise _LifecycleAbort("decode interrupted") + + +def test_query_baseexception_releases_prepared_and_reader_exactly_once(): + events = [] + reader = _PreparedReaderStub( + abort_query=True, + close_events=events, + ) + pool = _ReaderPoolStub(reader, events) + prepared_cache = runtime_module.PreparedSeriesCache( + _PreparedIndexStub(), max_entries=0 + ) + runtime = SimpleNamespace( + index=_PreparedIndexStub(), readers=pool, prepared=prepared_cache + ) + + with pytest.raises(_LifecycleAbort, match="query interrupted"): + RuntimeSeriesReader(runtime, 0)._query_at_locator(0, offset=0, limit=1) + + assert [item.close_calls for item in reader.prepared] == [1] + assert pool.acquire_calls == 1 + assert pool.release_calls == 1 + assert events == ["prepared.close:0", "reader.release"] + prepared_cache.close() + + +def test_decode_baseexception_closes_result_before_ownership_leases(): + events = [] + result = _FailingResultStub(events) + reader = _PreparedReaderStub( + close_events=events, + query_result=result, + ) + pool = _ReaderPoolStub(reader, events) + prepared_cache = runtime_module.PreparedSeriesCache( + _PreparedIndexStub(), max_entries=0 + ) + runtime = SimpleNamespace( + index=_PreparedIndexStub(), readers=pool, prepared=prepared_cache + ) + + with pytest.raises(_LifecycleAbort, match="decode interrupted"): + RuntimeSeriesReader(runtime, 0)._query_at_locator(0, offset=0, limit=1) + + assert result.close_calls == 1 + assert [item.close_calls for item in reader.prepared] == [1] + assert pool.release_calls == 1 + assert events == ["result.close", "prepared.close:0", "reader.release"] + prepared_cache.close() + + +def test_reader_pool_close_is_idempotent_and_closes_past_errors(monkeypatch): + pool = runtime_module.ReaderSessionPool( + SimpleNamespace(), max_open_files=2, validate_generation=False + ) + sessions = {} + + class _Session: + def __init__(self, file_id): + self.reader = file_id + self.active_uses = 0 + self.close_calls = 0 + + def close(self): + self.close_calls += 1 + if self.reader == 0: + raise RuntimeError("reader close 0") + + def new_session(file_id): + session = _Session(file_id) + sessions[file_id] = session + return session + + monkeypatch.setattr(pool, "_new_session", new_session) + with pool.acquire(0): + pass + with pool.acquire(1): + pass + + with pytest.raises(RuntimeError, match="reader close 0"): + pool.close() + assert sessions[0].close_calls == 1 + assert sessions[1].close_calls == 1 + pool.close() + assert sessions[0].close_calls == 1 + assert sessions[1].close_calls == 1 + + +def test_resultset_survives_zero_retention_prepared_eviction(tmp_path, monkeypatch): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + monkeypatch.setenv("TSFILE_DATAFRAME_MAX_PREPARED_SERIES", "0") + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + runtime = dataframe._runtime + series = runtime.index.record(LOGICAL_SERIES, 0) + span = runtime.index.record(SERIES_FILE_SPAN, series[2]) + with runtime.readers.acquire(0) as reader: + prepared_lease = runtime.prepared.acquire(0, span[2], reader) + prepared = prepared_lease.__enter__() + result = reader.query_prepared(prepared, offset=0, limit=2) + prepared_lease.__exit__(None, None, None) + assert runtime.prepared.size == 0 + timestamps, values = RuntimeSeriesReader._consume(result) + + np.testing.assert_array_equal(timestamps, np.array([0, 1], dtype=np.int64)) + np.testing.assert_array_equal(values, np.array([0.0, 1.0])) + + +def _write_runtime_aligned_fields_file(path): + schema = TableSchema( + "weather", + [ + ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), + ColumnSchema("value_a", TSDataType.DOUBLE, ColumnCategory.FIELD), + ColumnSchema("value_b", TSDataType.DOUBLE, ColumnCategory.FIELD), + ], + ) + with TsFileTableWriter(str(path), schema) as writer: + writer.write_dataframe( + pd.DataFrame( + { + "time": [0, 1, 2], + "device": ["d0", "d0", "d0"], + "value_a": [1.0, 2.0, 3.0], + "value_b": [10.0, 20.0, 30.0], + } + ) + ) + + +@pytest.mark.parametrize("release_owner_first", [True, False]) +def test_aligned_prepared_handles_survive_both_eviction_orders( + tmp_path, monkeypatch, release_owner_first +): + source = tmp_path / "aligned.tsfile" + _write_runtime_aligned_fields_file(source) + monkeypatch.setenv("TSFILE_DATAFRAME_MAX_PREPARED_SERIES", "0") + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + runtime = dataframe._runtime + series_records = [ + runtime.index.record(LOGICAL_SERIES, index) for index in range(2) + ] + spans = [ + runtime.index.record(SERIES_FILE_SPAN, series[2]) + for series in series_records + ] + with runtime.readers.acquire(0) as reader: + owner_lease = runtime.prepared.acquire(0, spans[0][2], reader) + owner = owner_lease.__enter__() + value_lease = runtime.prepared.acquire( + 0, spans[1][2], reader, time_owner=owner + ) + value = value_lease.__enter__() + + if release_owner_first: + owner_lease.__exit__(None, None, None) + result = reader.query_prepared(value, offset=0, limit=3) + value_lease.__exit__(None, None, None) + expected = np.array([10.0, 20.0, 30.0]) + else: + value_lease.__exit__(None, None, None) + result = reader.query_prepared(owner, offset=0, limit=3) + owner_lease.__exit__(None, None, None) + expected = np.array([1.0, 2.0, 3.0]) + + assert runtime.prepared.size == 0 + timestamps, values = RuntimeSeriesReader._consume(result) + + np.testing.assert_array_equal(timestamps, np.array([0, 1, 2], dtype=np.int64)) + np.testing.assert_array_equal(values, expected) + + +def test_dataframe_reopens_process_local_runtime_after_fork(tmp_path): + if not _linux_proc_resource_attribution_available(): + pytest.skip("Linux fork lifecycle validation is unavailable") + + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + parent_runtime = dataframe._runtime + parent_reader_pool = parent_runtime.readers + parent_prepared_cache = parent_runtime.prepared + parent_index = parent_runtime.index + parent_executor = parent_runtime._query_executor + inherited_executor_threads = () + if parent_executor is not None: + assert ( + parent_executor.submit(lambda: "parent-ready").result() + == "parent-ready" + ) + inherited_executor_threads = tuple(parent_executor._threads) + assert inherited_executor_threads + inherited_runtime_lease = dataframe._runtime_lease + series = dataframe[0] + np.testing.assert_array_equal(series[:], np.array([0.0, 1.0])) + series.close() + assert parent_reader_pool.open_count == 1 + assert parent_prepared_cache.size == 1 + + read_fd, write_fd = os.pipe() + child_pid = os.fork() + if child_pid == 0: + os.close(read_fd) + try: + child_series = dataframe[0] + child_runtime = dataframe._runtime + assert child_runtime is not parent_runtime + assert child_runtime.readers is not parent_reader_pool + assert child_runtime.prepared is not parent_prepared_cache + assert child_runtime.creator_pid == os.getpid() + assert parent_runtime._resources_closed + assert parent_runtime._discarded_pid == os.getpid() + assert parent_runtime.readers is None + assert parent_runtime.prepared is None + assert parent_runtime.index is None + assert parent_runtime.catalog is None + if parent_executor is not None: + assert parent_executor._shutdown + assert parent_executor._work_queue is None + assert parent_executor._threads == set() + assert all( + thread._target is None + and thread._args == () + and thread._kwargs == {} + for thread in inherited_executor_threads + ) + assert parent_reader_pool._closed + assert parent_reader_pool._sessions == {} + assert parent_prepared_cache._closed + assert parent_prepared_cache._entries == {} + assert parent_index._view is None + assert inherited_runtime_lease._closed + values = child_series[:] + child_series.close() + np.testing.assert_array_equal(values, np.array([0.0, 1.0])) + dataframe.close() + assert child_runtime._resources_closed + if os.path.isdir("/proc/self/fd"): + open_targets = { + os.path.realpath(f"/proc/self/fd/{fd}") + for fd in os.listdir("/proc/self/fd") + if os.path.exists(f"/proc/self/fd/{fd}") + } + assert os.path.realpath(source) not in open_targets + assert os.path.realpath(parent_index.path) not in open_targets + os.write(write_fd, b"OK") + status = 0 + except BaseException as exc: + os.write(write_fd, f"{type(exc).__name__}: {exc}".encode()) + status = 1 + finally: + os.close(write_fd) + os._exit(status) + + os.close(write_fd) + child_message = os.read(read_fd, 4096) + os.close(read_fd) + waited_pid, child_status = os.waitpid(child_pid, 0) + + assert waited_pid == child_pid + assert os.WIFEXITED(child_status), child_message.decode() + assert os.WEXITSTATUS(child_status) == 0, child_message.decode() + assert child_message == b"OK" + assert dataframe._runtime is parent_runtime + assert dataframe._runtime.readers is parent_reader_pool + assert dataframe._runtime.prepared is parent_prepared_cache + assert parent_reader_pool.open_count == 1 + assert parent_prepared_cache.size == 1 + if parent_executor is not None: + assert not parent_executor._shutdown + assert parent_executor._work_queue is not None + assert parent_executor.submit(lambda: "parent-ok").result() == "parent-ok" + parent_series = dataframe[0] + np.testing.assert_array_equal(parent_series[:], np.array([0.0, 1.0])) + parent_series.close() + + +def test_prefork_timeseries_fails_fast_in_child_without_harming_parent(tmp_path): + if not _linux_proc_resource_attribution_available(): + pytest.skip("Linux fork lifecycle validation is unavailable") + + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + parent_runtime = dataframe._runtime + parent_reader_pool = parent_runtime.readers + parent_prepared_cache = parent_runtime.prepared + parent_index = parent_runtime.index + series = dataframe[0] + np.testing.assert_array_equal(series[:], np.array([0.0, 1.0])) + inherited_query_lease = series._runtime_lease.query_lease() + + read_fd, write_fd = os.pipe() + child_pid = os.fork() + if child_pid == 0: + os.close(read_fd) + try: + with pytest.raises( + RuntimeError, + match="obtain a new Timeseries from the child.*TsFileDataFrame", + ): + series[:] + with pytest.raises( + RuntimeError, + match="obtain a new Timeseries from the child.*TsFileDataFrame", + ): + series._runtime_lease.clone() + inherited_query_lease.close() + assert not inherited_query_lease._acquired + series.close() + assert series._closed + assert parent_runtime._resources_closed + assert parent_runtime._discarded_pid == os.getpid() + assert parent_runtime.readers is None + assert parent_runtime.prepared is None + assert parent_runtime.index is None + assert parent_runtime.catalog is None + assert parent_reader_pool._closed + assert parent_reader_pool._sessions == {} + assert parent_prepared_cache._closed + assert parent_prepared_cache._entries == {} + assert parent_index._view is None + os.write(write_fd, b"OK") + status = 0 + except BaseException as exc: + os.write(write_fd, f"{type(exc).__name__}: {exc}".encode()) + status = 1 + finally: + os.close(write_fd) + os._exit(status) + + os.close(write_fd) + child_message = os.read(read_fd, 4096) + os.close(read_fd) + waited_pid, child_status = os.waitpid(child_pid, 0) + + assert waited_pid == child_pid + assert os.WIFEXITED(child_status), child_message.decode() + assert os.WEXITSTATUS(child_status) == 0, child_message.decode() + assert child_message == b"OK" + assert inherited_query_lease._acquired + assert parent_runtime._query_leases == 1 + np.testing.assert_array_equal(series[:], np.array([0.0, 1.0])) + inherited_query_lease.close() + assert parent_runtime._query_leases == 0 + series.close() diff --git a/python/tsfile/dataset/dataframe.py b/python/tsfile/dataset/dataframe.py index 06f7109d2..c81990777 100644 --- a/python/tsfile/dataset/dataframe.py +++ b/python/tsfile/dataset/dataframe.py @@ -835,6 +835,7 @@ def _from_subset( cls, parent: "TsFileDataFrame", series_refs: List[SeriesRefKey] ) -> "TsFileDataFrame": """Create a lightweight view that reuses the parent's readers and caches.""" + parent._assert_open() obj = object.__new__(cls) obj._root = parent._root if parent._is_view else parent obj._is_view = True @@ -864,9 +865,49 @@ def _from_subset( def _owner(self) -> "TsFileDataFrame": return self + def _ensure_process_local_runtime(self): + runtime = self._runtime + if runtime is None or runtime.creator_pid == os.getpid(): + return + + subset_refs = list(self._index.series) if self._is_view else None + inherited_lease = self._runtime_lease + replacement = None + if self._is_view and self._root is not None and not self._root._closed: + self._root._ensure_process_local_runtime() + replacement = self._root._runtime + if replacement is None: + replacement = runtime.fork_replacement() + + replacement_lease = replacement.lease() + try: + if inherited_lease is None: + runtime.discard_after_fork() + else: + inherited_lease.close() + except BaseException: + replacement_lease.close() + raise + self._runtime = replacement + self._runtime_lease = replacement_lease + catalog = replacement.catalog + if subset_refs is None: + self._index = catalog + else: + self._index = SimpleNamespace( + model=catalog.model, + table_entries=catalog.table_entries, + devices=catalog.devices, + device_index=catalog.device_index, + device_time_bounds=catalog.device_time_bounds, + series=subset_refs, + series_shards=catalog.series_shards, + ) + def _assert_open(self): if self._closed: raise RuntimeError("Current TsFileDataFrame is closed.") + self._ensure_process_local_runtime() @contextlib.contextmanager def _query_guard(self): @@ -1269,6 +1310,7 @@ def _get_timeseries( ) def __getitem__(self, key): + self._assert_open() try: import pandas as pd @@ -1457,7 +1499,9 @@ def close(self): return self._closed = True if self._runtime_lease is not None: - self._runtime_lease.close() + runtime_lease = self._runtime_lease + self._runtime_lease = None + runtime_lease.close() else: for reader in self._readers.values(): reader.close() diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py index 1202123e4..7d62f66c3 100644 --- a/python/tsfile/dataset/runtime.py +++ b/python/tsfile/dataset/runtime.py @@ -23,6 +23,7 @@ from collections.abc import Mapping, Sequence import contextlib from concurrent.futures import ThreadPoolExecutor, wait +from concurrent.futures import thread as thread_pool_module from dataclasses import dataclass import os import threading @@ -58,6 +59,62 @@ _SERIES_DESCRIPTOR_CACHE_SIZE = 4096 +def _configured_non_negative_int(name, default): + raw = os.environ.get(name) + if raw is None: + return default + try: + value = int(raw) + except ValueError as exc: + raise ValueError(f"{name} must be an integer") from exc + if value < 0: + raise ValueError(f"{name} must be non-negative") + return value + + +def _configured_prepared_series_cache_size(): + return _configured_non_negative_int( + "TSFILE_DATAFRAME_MAX_PREPARED_SERIES", + _SERIES_DESCRIPTOR_CACHE_SIZE, + ) + + +def _run_cleanups(cleanups): + first_error = None + for cleanup in cleanups: + try: + cleanup() + except BaseException as exc: + if first_error is None: + first_error = exc + if first_error is not None: + raise first_error + + +def _discard_executor_after_fork(executor): + """Detach a CPython ThreadPoolExecutor without touching inherited locks.""" + # CPython resets the thread module's global shutdown lock after fork, but + # not a ThreadPoolExecutor instance's _shutdown_lock. shutdown()/join() + # can therefore deadlock in the child. Disconnect the vanished workers + # and their queues lock-free; this runtime never uses the executor again. + executor._shutdown = True + threads = tuple(getattr(executor, "_threads", ())) + if hasattr(executor, "_threads"): + executor._threads = set() + if hasattr(executor, "_work_queue"): + executor._work_queue = None + thread_queues = getattr(thread_pool_module, "_threads_queues", None) + for thread in threads: + if thread_queues is not None: + thread_queues.pop(thread, None) + if hasattr(thread, "_target"): + thread._target = None + if hasattr(thread, "_args"): + thread._args = () + if hasattr(thread, "_kwargs"): + thread._kwargs = {} + + @dataclass(frozen=True) class RuntimeSeriesShard: """One immutable physical fragment already expanded from the mmap route.""" @@ -102,17 +159,27 @@ def _exact_tag_filter(tag_columns, tag_values): class RuntimeLease: def __init__(self, runtime: "DatasetRuntime"): self._runtime = runtime + self._creator_pid = os.getpid() self._lock = threading.Lock() self._closed = False runtime._acquire_object() + def _assert_current_process(self): + if self._creator_pid != os.getpid(): + raise RuntimeError( + "Inherited runtime handles cannot be queried after fork; " + "obtain a new Timeseries from the child's TsFileDataFrame" + ) + def clone(self): + self._assert_current_process() with self._lock: if self._closed: raise RuntimeError("Runtime lease is closed") return RuntimeLease(self._runtime) def query_lease(self): + self._assert_current_process() with self._lock: if self._closed: raise RuntimeError("Runtime lease is closed") @@ -120,6 +187,12 @@ def query_lease(self): return _QueryLease(self._runtime, acquired=True) def close(self): + if self._creator_pid != os.getpid(): + if self._closed: + return + self._closed = True + self._runtime.discard_after_fork() + return with self._lock: if self._closed: return @@ -130,9 +203,15 @@ def close(self): class _QueryLease: def __init__(self, runtime: "DatasetRuntime", acquired: bool = False): self._runtime = runtime + self._creator_pid = os.getpid() self._acquired = acquired def __enter__(self): + if self._creator_pid != os.getpid(): + raise RuntimeError( + "Inherited runtime handles cannot be queried after fork; " + "obtain a new Timeseries from the child's TsFileDataFrame" + ) if not self._acquired: self._runtime._acquire_query() self._acquired = True @@ -142,6 +221,12 @@ def __exit__(self, *_): self.close() def close(self): + if self._creator_pid != os.getpid(): + if not self._acquired: + return + self._acquired = False + self._runtime.discard_after_fork() + return if self._acquired: self._acquired = False self._runtime._release_query() @@ -165,6 +250,7 @@ def __init__( self._validate_generation() self.reader = TsFileReaderPy(path) self.active_uses = 0 + self._closed = False def _validate_generation(self): st = os.stat(self.path) @@ -178,6 +264,9 @@ def _validate_generation(self): ) def close(self): + if self._closed: + return + self._closed = True self.reader.close() @@ -256,8 +345,16 @@ def close(self): self._condition.wait() sessions = list(self._sessions.values()) self._sessions.clear() - for session in sessions: - session.close() + _run_cleanups(session.close for session in sessions) + + def discard_after_fork(self): + """Close only this child copy without acquiring inherited locks.""" + if self._closed and not self._sessions: + return + self._closed = True + sessions = list(self._sessions.values()) + self._sessions = OrderedDict() + _run_cleanups(session.close for session in sessions) @property def open_count(self): @@ -265,18 +362,32 @@ def open_count(self): return len(self._sessions) +@dataclass +class _PreparedSeriesEntry: + prepared: object + active_uses: int = 0 + time_owner: Optional["_PreparedSeriesEntry"] = None + dependent_uses: int = 0 + + class PreparedSeriesCache: - """Runtime-wide single-flight cache of native exact-locator metadata.""" + """Runtime-wide single-flight LRU of native exact-locator metadata.""" def __init__( self, index: MappedDatasetIndex, validate_references: bool = True, + max_entries: Optional[int] = None, ): self._index = index self._validate_references = bool(validate_references) + maximum = _SERIES_DESCRIPTOR_CACHE_SIZE if max_entries is None else max_entries + if int(maximum) < 0: + raise ValueError("max_entries must be non-negative") + self.max_entries = int(maximum) self._condition = threading.Condition() - self._entries = {} + self._entries = OrderedDict() + self._entries_by_prepared_id = {} self._loading = set() self._closed = False @@ -300,48 +411,165 @@ def _locator_tuple(self, file_id, locator_id): device_span[3], ) - def get(self, file_id, locator_id, reader, time_owner=None): + def _evict_idle_locked(self): + evicted = [] + while len(self._entries) > self.max_entries: + idle_key = next( + ( + key + for key, entry in self._entries.items() + if entry.active_uses == 0 and entry.dependent_uses == 0 + ), + None, + ) + if idle_key is None: + break + evicted.append(self._pop_entry_locked(idle_key)) + return evicted + + def _pop_entry_locked(self, key): + entry = self._entries.pop(key) + self._entries_by_prepared_id.pop(id(entry.prepared), None) + if entry.time_owner is not None: + entry.time_owner.dependent_uses -= 1 + entry.time_owner = None + return entry.prepared + + def _drain_entries_locked(self): + prepared = [] + while self._entries: + leaf_key = next( + ( + key + for key, entry in self._entries.items() + if entry.dependent_uses == 0 + ), + None, + ) + if leaf_key is None: + raise RuntimeError("PreparedSeries owner graph contains a cycle") + prepared.append(self._pop_entry_locked(leaf_key)) + return prepared + + @staticmethod + def _close_entries(entries): + _run_cleanups(prepared.close for prepared in entries) + + def _release(self, entry): + with self._condition: + entry.active_uses -= 1 + evicted = self._evict_idle_locked() + self._condition.notify_all() + self._close_entries(evicted) + + @contextlib.contextmanager + def acquire(self, file_id, locator_id, reader, time_owner=None): key = (id(self._index), file_id, locator_id) + entry = None + owner_entry = None + prepared_time_owner = time_owner with self._condition: while True: if self._closed: raise RuntimeError("PreparedSeriesCache is closed") - result = self._entries.get(key) - if result is not None: - return result + entry = self._entries.get(key) + if entry is not None: + self._entries.move_to_end(key) + entry.active_uses += 1 + break if key not in self._loading: + if time_owner is not None: + requested_owner = self._entries_by_prepared_id.get( + id(time_owner) + ) + if ( + requested_owner is None + or requested_owner.prepared is not time_owner + or requested_owner.active_uses == 0 + ): + raise RuntimeError( + "PreparedSeries time owner requires an active cache lease" + ) + owner_entry = requested_owner.time_owner or requested_owner + prepared_time_owner = owner_entry.prepared self._loading.add(key) break self._condition.wait() - try: - result = reader.prepare_series( - self._locator_tuple(file_id, locator_id), time_owner=time_owner - ) - except Exception: + + evicted = [] + if entry is None: + try: + result = reader.prepare_series( + self._locator_tuple(file_id, locator_id), + time_owner=prepared_time_owner, + ) + except BaseException: + with self._condition: + self._loading.discard(key) + self._condition.notify_all() + raise + with self._condition: - self._loading.remove(key) - self._condition.notify_all() - raise - with self._condition: - if self._closed: - result.close() - self._loading.remove(key) + self._loading.discard(key) + rejected = self._closed + if not rejected: + entry = _PreparedSeriesEntry( + result, + active_uses=1, + time_owner=owner_entry, + ) + if owner_entry is not None: + owner_entry.dependent_uses += 1 + self._entries[key] = entry + self._entries_by_prepared_id[id(result)] = entry + evicted = self._evict_idle_locked() self._condition.notify_all() + if rejected: + try: + result.close() + except BaseException as exc: + raise RuntimeError("PreparedSeriesCache is closed") from exc raise RuntimeError("PreparedSeriesCache is closed") - self._entries[key] = result - self._loading.remove(key) - self._condition.notify_all() - return result + + try: + self._close_entries(evicted) + except BaseException: + try: + self._release(entry) + except BaseException: + pass + raise + + try: + yield entry.prepared + except BaseException: + try: + self._release(entry) + except BaseException: + pass + raise + else: + self._release(entry) def close(self): with self._condition: self._closed = True - while self._loading: + while self._loading or any( + entry.active_uses for entry in self._entries.values() + ): self._condition.wait() - entries = list(self._entries.values()) - self._entries.clear() - for prepared in entries: - prepared.close() + entries = self._drain_entries_locked() + self._close_entries(entries) + + def discard_after_fork(self): + """Close only this child copy without acquiring inherited locks.""" + if self._closed and not self._entries and not self._loading: + return + self._closed = True + entries = self._drain_entries_locked() + self._entries_by_prepared_id = {} + self._loading = set() + self._close_entries(entries) @property def size(self): @@ -359,12 +587,12 @@ def __init__( trust_index: bool = False, ): self.trust_index = bool(trust_index) - self.index = MappedDatasetIndex(path, trust_index=self.trust_index) maximum = ( int(os.environ.get("TSFILE_DATAFRAME_MAX_OPEN_FILES", "16")) if max_open_files is None else max_open_files ) + self.max_open_files = max(1, int(maximum)) workers = ( int( os.environ.get( @@ -382,38 +610,139 @@ def __init__( else query_parallel_min_rows ) self.query_parallel_min_rows = max(1, int(minimum_rows)) - self._query_executor = ( - ThreadPoolExecutor( - max_workers=self.query_workers, - thread_name_prefix="tsfile-dataframe-query", - ) - if self.query_workers > 1 - else None - ) - self.readers = ReaderSessionPool( - self.index, - maximum, - validate_generation=not self.trust_index, - ) - self.prepared = PreparedSeriesCache( - self.index, - validate_references=not self.trust_index, - ) + prepared_max_entries = _configured_prepared_series_cache_size() + self.creator_pid = os.getpid() + self._discarded_pid = None + self._path = os.fspath(path) + self.index = None + self._query_executor = None + self.readers = None + self.prepared = None + self.catalog = None + self._resources_closed = False self._condition = threading.Condition() self._object_leases = 0 self._query_leases = 0 self._accepting = True self._torn_down = False - self.catalog = MappedDataFrameCatalog(self) + + try: + self.index = MappedDatasetIndex(self._path, trust_index=self.trust_index) + self._query_executor = ( + ThreadPoolExecutor( + max_workers=self.query_workers, + thread_name_prefix="tsfile-dataframe-query", + ) + if self.query_workers > 1 + else None + ) + self.readers = ReaderSessionPool( + self.index, + self.max_open_files, + validate_generation=not self.trust_index, + ) + self.prepared = PreparedSeriesCache( + self.index, + validate_references=not self.trust_index, + max_entries=prepared_max_entries, + ) + self.catalog = MappedDataFrameCatalog(self) + except BaseException: + try: + self._close_resources() + except BaseException: + pass + raise + + def _close_resources(self): + if self._resources_closed: + return + self._resources_closed = True + cleanups = [] + if self._query_executor is not None: + cleanups.append( + lambda executor=self._query_executor: executor.shutdown( + wait=True, cancel_futures=True + ) + ) + if self.prepared is not None: + cleanups.append(self.prepared.close) + if self.readers is not None: + cleanups.append(self.readers.close) + if self.index is not None: + cleanups.append(self.index.close) + _run_cleanups(cleanups) + + def fork_replacement(self): + """Create fresh process-local mutable/native state after ``fork``.""" + return type(self)( + self._path, + max_open_files=self.max_open_files, + query_workers=self.query_workers, + query_parallel_min_rows=self.query_parallel_min_rows, + trust_index=self.trust_index, + ) + + def _assert_current_process(self): + if self.creator_pid != os.getpid(): + raise RuntimeError( + "Inherited DatasetRuntime cannot be used after fork; " + "obtain a new Timeseries from the child's TsFileDataFrame" + ) + + def discard_after_fork(self): + """Deterministically discard this child copy without inherited locks.""" + current_pid = os.getpid() + if self.creator_pid == current_pid: + raise RuntimeError("discard_after_fork requires an inherited runtime") + if self._discarded_pid == current_pid: + return + + self._discarded_pid = current_pid + self._resources_closed = True + self._accepting = False + self._torn_down = True + + prepared = self.prepared + readers = self.readers + index = self.index + catalog = self.catalog + executor = self._query_executor + self._query_executor = None + self.prepared = None + self.readers = None + self.index = None + self.catalog = None + + if catalog is not None: + catalog.runtime = None + catalog.index = None + catalog_readers = getattr(catalog, "_readers", None) + if catalog_readers is not None: + catalog_readers.clear() + + cleanups = [] + if executor is not None: + cleanups.append(lambda: _discard_executor_after_fork(executor)) + if prepared is not None: + cleanups.append(prepared.discard_after_fork) + if readers is not None: + cleanups.append(readers.discard_after_fork) + if index is not None: + cleanups.append(index.close) + _run_cleanups(cleanups) def lease(self): + self._assert_current_process() return RuntimeLease(self) def query_lease(self): + self._assert_current_process() return _QueryLease(self) def map_query_groups(self, function, groups, estimated_rows=None): """Run independent query groups under the caller's query lease.""" + self._assert_current_process() groups = list(groups) if not groups: return [] @@ -435,12 +764,16 @@ def map_query_groups(self, function, groups, estimated_rows=None): raise def _acquire_object(self): + self._assert_current_process() with self._condition: if not self._accepting: raise RuntimeError("Dataset Runtime is closing") self._object_leases += 1 def _release_object(self): + if self.creator_pid != os.getpid(): + self.discard_after_fork() + return teardown = False with self._condition: self._object_leases -= 1 @@ -451,19 +784,19 @@ def _release_object(self): teardown = not self._torn_down self._torn_down = True if teardown: - if self._query_executor is not None: - self._query_executor.shutdown(wait=True, cancel_futures=True) - self.prepared.close() - self.readers.close() - self.index.close() + self._close_resources() def _acquire_query(self): + self._assert_current_process() with self._condition: if not self._accepting: raise RuntimeError("Dataset Runtime is closing") self._query_leases += 1 def _release_query(self): + if self.creator_pid != os.getpid(): + self.discard_after_fork() + return with self._condition: self._query_leases -= 1 self._condition.notify_all() @@ -960,14 +1293,16 @@ def _query_at_locator( limit=None, ): with self.runtime.readers.acquire(self.file_id) as reader: - prepared = self.runtime.prepared.get(self.file_id, locator_id, reader) - if offset is None: - result = reader.query_prepared( - prepared, start_time=start_time, end_time=end_time - ) - else: - result = reader.query_prepared(prepared, offset=offset, limit=limit) - return self._consume(result) + with self.runtime.prepared.acquire( + self.file_id, locator_id, reader + ) as prepared: + if offset is None: + result = reader.query_prepared( + prepared, start_time=start_time, end_time=end_time + ) + else: + result = reader.query_prepared(prepared, offset=offset, limit=limit) + return self._consume(result) def read_series_by_ref(self, device_id, column_id, start_time, end_time): return self._query(device_id, column_id, start_time, end_time) @@ -1010,19 +1345,25 @@ def read_device_fields_by_time_range( self._identity(device_id, column_id)[3] for column_id in column_ids ] with self.runtime.readers.acquire(self.file_id) as reader: - prepared = [] - time_owner = None - for span in spans: - current = self.runtime.prepared.get( - self.file_id, span[2], reader, time_owner=time_owner + with contextlib.ExitStack() as stack: + prepared = [] + time_owner = None + for span in spans: + current = stack.enter_context( + self.runtime.prepared.acquire( + self.file_id, + span[2], + reader, + time_owner=time_owner, + ) + ) + prepared.append(current) + if time_owner is None: + time_owner = current + result = reader.query_prepared_multi( + prepared, start_time=start_time, end_time=end_time ) - prepared.append(current) - if time_owner is None: - time_owner = current - result = reader.query_prepared_multi( - prepared, start_time=start_time, end_time=end_time - ) - return self._consume_multi(result, column_names) + return self._consume_multi(result, column_names) # Non-aligned device (or fields spanning different device spans): each # field carries its own timeline, so align them onto a single timestamp From 94649f3e4a1d0f1245cf0616ac24c3853db2769a Mon Sep 17 00:00:00 2001 From: Gewu <89496957+RkGrit@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:31:37 +0800 Subject: [PATCH 3/3] feat(python): add trusted index and data attestations --- python/tests/test_dataset_index.py | 202 +++++++++++++++++++ python/tests/test_index_identity.py | 213 ++++++++++++++++++++ python/tsfile/dataset/__init__.py | 27 ++- python/tsfile/dataset/dataframe.py | 64 +++++- python/tsfile/dataset/index_identity.py | 254 ++++++++++++++++++++++++ python/tsfile/dataset/runtime.py | 142 ++++++++++++- 6 files changed, 890 insertions(+), 12 deletions(-) create mode 100644 python/tests/test_index_identity.py create mode 100644 python/tsfile/dataset/index_identity.py diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index 9b3eab077..684ff2f3c 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -24,6 +24,8 @@ import pyarrow as pa import pytest +import tsfile.dataset.dataframe as dataframe_module +import tsfile.dataset.index_identity as identity_module import tsfile.dataset.index as index_module import tsfile.dataset.runtime as runtime_module from tsfile import ( @@ -1746,3 +1748,203 @@ def test_prefork_timeseries_fails_fast_in_child_without_harming_parent(tmp_path) inherited_query_lease.close() assert parent_runtime._query_leases == 0 series.close() + + +def _build_runtime_attestations(paths): + canonical_paths = [os.path.realpath(os.fspath(path)) for path in paths] + with TsFileDataFrame( + canonical_paths, + show_progress=False, + use_index=True, + ) as dataframe: + index_path = dataframe._runtime.index.path + + index_attestation = identity_module.attest_index(index_path) + with MappedDatasetIndex(index_path, trust_index=True) as mapped_index: + data_attestations = tuple( + identity_module.attest_data_file( + file_id, + mapped_index.string(record[0]), + existing_index_fingerprint=f"{record[3]:016x}", + ) + for file_id in range(mapped_index.count(TSFILE_RECORD)) + for record in (mapped_index.record(TSFILE_RECORD, file_id),) + ) + return index_path, index_attestation, data_attestations + + +def test_attested_trusted_dataframe_uses_manifest_without_directory_scan( + tmp_path, monkeypatch +): + paths = [tmp_path / "part-0.tsfile", tmp_path / "part-1.tsfile"] + _write_runtime_file(paths[0], 0) + _write_runtime_file(paths[1], 10) + _, index_attestation, data_attestations = _build_runtime_attestations(paths) + + def unexpected_scan(_paths): + raise AssertionError("attested workers must not expand or scan dataset paths") + + monkeypatch.setattr(dataframe_module, "_expand_paths", unexpected_scan) + with TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + index_attestation=index_attestation, + data_file_attestations=data_attestations, + ) as dataframe: + runtime = dataframe._runtime + assert tuple(dataframe._paths) == tuple( + item.canonical_path for item in data_attestations + ) + assert runtime.data_manifest_identity_sha256 == ( + identity_module.data_manifest_identity( + data_attestations, + require_sorted=True, + ) + ) + assert tuple(runtime.file_generation_tokens) == (0, 1) + assert runtime.prepared._mapped_index_identity == int.from_bytes( + bytes.fromhex(index_attestation.index_identity_sha256)[:8], + byteorder="little", + ) + np.testing.assert_array_equal( + dataframe[0][:], + np.array([0.0, 1.0, 10.0, 11.0]), + ) + + +def test_attested_dataframe_requires_explicit_trust_and_complete_pair(tmp_path): + path = tmp_path / "part.tsfile" + _write_runtime_file(path, 0) + _, index_attestation, data_attestations = _build_runtime_attestations((path,)) + + with pytest.raises(ValueError, match="trust_index=True"): + TsFileDataFrame( + str(tmp_path), + show_progress=False, + use_index=True, + index_attestation=index_attestation, + data_file_attestations=data_attestations, + ) + with pytest.raises(ValueError, match="provided together"): + TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + index_attestation=index_attestation, + ) + with pytest.raises(ValueError, match="provided together"): + TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + data_file_attestations=data_attestations, + ) + + +def test_attested_dataframe_rejects_incomplete_order_and_unbound_fingerprint(tmp_path): + paths = [tmp_path / "part-0.tsfile", tmp_path / "part-1.tsfile"] + _write_runtime_file(paths[0], 0) + _write_runtime_file(paths[1], 10) + _, index_attestation, data_attestations = _build_runtime_attestations(paths) + + with pytest.raises(ValueError, match="file_id=0..N-1 in order"): + TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + index_attestation=index_attestation, + data_file_attestations=data_attestations[:1], + ) + with pytest.raises(ValueError, match="file_id=0..N-1 in order"): + TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + index_attestation=index_attestation, + data_file_attestations=tuple(reversed(data_attestations)), + ) + + first = data_attestations[0] + unbound = identity_module.DataFileAttestation( + file_id=first.file_id, + canonical_path=first.canonical_path, + st_dev=first.st_dev, + st_ino=first.st_ino, + st_size=first.st_size, + st_mtime_ns=first.st_mtime_ns, + ) + with pytest.raises( + identity_module.AttestationMismatchError, + match="fingerprint", + ): + TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + index_attestation=index_attestation, + data_file_attestations=(unbound, data_attestations[1]), + ) + + +def test_attested_index_mismatch_fails_before_mapping(tmp_path, monkeypatch): + path = tmp_path / "part.tsfile" + _write_runtime_file(path, 0) + index_path, index_attestation, data_attestations = _build_runtime_attestations( + (path,) + ) + with open(index_path, "ab") as stream: + stream.write(b"x") + + mapped_calls = 0 + + def forbidden_mapping(*_args, **_kwargs): + nonlocal mapped_calls + mapped_calls += 1 + raise AssertionError("a mismatched index must not be mapped") + + monkeypatch.setattr(runtime_module, "MappedDatasetIndex", forbidden_mapping) + with pytest.raises( + identity_module.AttestationMismatchError, + match="index stat", + ): + TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + index_attestation=index_attestation, + data_file_attestations=data_attestations, + ) + assert mapped_calls == 0 + + +def test_attested_data_mismatch_fails_before_native_reader_open(tmp_path, monkeypatch): + path = tmp_path / "part.tsfile" + _write_runtime_file(path, 0) + _, index_attestation, data_attestations = _build_runtime_attestations((path,)) + + replacement = tmp_path / "replacement.tsfile" + replacement.write_bytes(b"X" * path.stat().st_size) + os.replace(replacement, path) + + native_open_calls = 0 + + def forbidden_reader(*_args, **_kwargs): + nonlocal native_open_calls + native_open_calls += 1 + raise AssertionError("a mismatched data file must not reach the native reader") + + monkeypatch.setattr(runtime_module, "TsFileReaderPy", forbidden_reader) + with TsFileDataFrame( + str(tmp_path), + show_progress=False, + trust_index=True, + index_attestation=index_attestation, + data_file_attestations=data_attestations, + ) as dataframe: + with pytest.raises( + identity_module.AttestationMismatchError, + match="data-file stat", + ): + dataframe[0][:] + assert native_open_calls == 0 diff --git a/python/tests/test_index_identity.py b/python/tests/test_index_identity.py new file mode 100644 index 000000000..ea3cd0ef4 --- /dev/null +++ b/python/tests/test_index_identity.py @@ -0,0 +1,213 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import dataclasses +import hashlib +import os +from pathlib import Path +import runpy +from types import SimpleNamespace + +import pytest + +import tsfile.dataset.index_identity as identity + + +def test_index_identity_module_imports_with_python39_dataclass_api(monkeypatch): + """The public dataset package supports Python 3.9, which has no slots= API.""" + + real_dataclass = dataclasses.dataclass + + def python39_dataclass(*args, **kwargs): + if "slots" in kwargs: + raise TypeError("dataclass() got an unexpected keyword argument 'slots'") + return real_dataclass(*args, **kwargs) + + monkeypatch.setattr(dataclasses, "dataclass", python39_dataclass) + runpy.run_path(identity.__file__, run_name="_tsfile_index_identity_py39_probe") + + +def test_index_attestation_hashes_exact_bytes_and_worker_validation_is_metadata_only( + tmp_path, monkeypatch +): + first = tmp_path / "first.tsidx" + first.write_bytes(b"final-index\x00bytes") + attestation = identity.attest_index(first, chunk_size=3) + assert ( + attestation.index_identity_sha256 + == hashlib.sha256(first.read_bytes()).hexdigest() + ) + + relocated = tmp_path / "relocated.tsidx" + relocated.write_bytes(first.read_bytes()) + relocated_attestation = identity.attest_index(relocated) + assert relocated_attestation.index_identity_sha256 == ( + attestation.index_identity_sha256 + ) + relocated.write_bytes(b"Final-index\x00bytes") + assert identity.attest_index(relocated).index_identity_sha256 != ( + attestation.index_identity_sha256 + ) + + hash_calls = 0 + + def forbidden_hash(*_args, **_kwargs): + nonlocal hash_calls + hash_calls += 1 + raise AssertionError("worker validation must not hash index contents") + + monkeypatch.setattr(identity, "attest_index", forbidden_hash) + for _ in range(32): + identity.validate_index_attestation(first, attestation) + assert hash_calls == 0 + + first.write_bytes(first.read_bytes() + b"x") + with pytest.raises(identity.AttestationMismatchError, match="index stat"): + identity.validate_index_attestation(first, attestation) + + +def test_index_attestation_fails_if_open_file_changes_while_hashing( + tmp_path, monkeypatch +): + path = tmp_path / "changing.tsidx" + path.write_bytes(b"stable-index-bytes") + real_fstat = identity.os.fstat + fstat_calls = 0 + + def changing_fstat(file_descriptor): + nonlocal fstat_calls + stat_result = real_fstat(file_descriptor) + fstat_calls += 1 + if fstat_calls == 1: + return stat_result + return SimpleNamespace( + st_dev=stat_result.st_dev, + st_ino=stat_result.st_ino, + st_size=stat_result.st_size, + st_mtime_ns=stat_result.st_mtime_ns + 1, + ) + + monkeypatch.setattr(identity.os, "fstat", changing_fstat) + with pytest.raises( + identity.AttestationMismatchError, match="changed while hashing" + ): + identity.attest_index(path, chunk_size=3) + assert fstat_calls == 2 + + +def test_manifest_and_generation_are_ordered_versioned_and_fail_closed(tmp_path): + paths = [tmp_path / "a.tsfile", tmp_path / "b.tsfile"] + for ordinal, path in enumerate(paths): + path.write_bytes(bytes([ordinal]) * (ordinal + 1)) + files = tuple(identity.attest_data_file(i, path) for i, path in enumerate(paths)) + + manifest = identity.data_manifest_identity(files) + assert len(manifest) == 64 and manifest == manifest.lower() + assert identity.data_manifest_identity(tuple(reversed(files))) != manifest + with pytest.raises(ValueError, match="strictly increasing"): + identity.data_manifest_identity(tuple(reversed(files)), require_sorted=True) + + token = identity.file_generation_token("11" * 32, manifest, files[0]) + assert token == identity.file_generation_token("11" * 32, manifest, files[0]) + assert token != identity.file_generation_token("11" * 32, manifest, files[1]) + identity.validate_data_file_attestation(paths[0], files[0]) + replacement = tmp_path / "replacement" + replacement.write_bytes(b"replacement") + os.replace(replacement, paths[0]) + with pytest.raises(identity.AttestationMismatchError, match="data-file stat"): + identity.validate_data_file_attestation(paths[0], files[0]) + + +@pytest.mark.parametrize( + "fingerprint", + ["", "0" * 15, "0" * 17, "ABCDEF0123456789", "0123456789abcdeg"], +) +def test_data_file_attestation_rejects_noncanonical_index_fingerprint( + tmp_path, fingerprint +): + path = tmp_path / "data.tsfile" + path.write_bytes(b"data") + with pytest.raises(ValueError, match="16-character lowercase hex"): + identity.attest_data_file(0, path, existing_index_fingerprint=fingerprint) + + +def test_data_file_fingerprint_has_unique_canonical_identity(tmp_path): + path = tmp_path / "data.tsfile" + path.write_bytes(b"data") + without_fingerprint = identity.attest_data_file(0, path) + with_fingerprint = identity.attest_data_file( + 0, path, existing_index_fingerprint="0123456789abcdef" + ) + manifest_without = identity.data_manifest_identity((without_fingerprint,)) + manifest_with = identity.data_manifest_identity((with_fingerprint,)) + assert manifest_without != manifest_with + assert identity.file_generation_token( + "11" * 32, manifest_without, without_fingerprint + ) != identity.file_generation_token("11" * 32, manifest_with, with_fingerprint) + + +def test_data_file_attestation_stats_the_opened_object(tmp_path, monkeypatch): + path = tmp_path / "data.tsfile" + path.write_bytes(b"data") + real_fstat = identity.os.fstat + fstat_calls = 0 + + def recording_fstat(file_descriptor): + nonlocal fstat_calls + fstat_calls += 1 + return real_fstat(file_descriptor) + + monkeypatch.setattr(identity.os, "fstat", recording_fstat) + attestation = identity.attest_data_file(7, path) + assert attestation.file_id == 7 + assert attestation.st_size == 4 + assert fstat_calls == 1 + + +@pytest.mark.parametrize("value", ["", "A" * 64, "0" * 63, "gg" * 32]) +def test_generation_token_digest_validation_is_strict(value): + with pytest.raises(ValueError, match="lowercase SHA-256"): + identity.file_generation_token(value, "44" * 32, object()) + + +def test_attestations_are_immutable_and_canonical_paths_are_absolute(tmp_path): + path = tmp_path / "index.tsidx" + path.write_bytes(b"index") + attestation = identity.attest_index(path) + assert Path(attestation.canonical_index_path).is_absolute() + with pytest.raises((AttributeError, TypeError)): + attestation.index_size = 1 + + +def test_dataset_identity_public_api_is_generic_and_explicit(): + import tsfile.dataset as dataset_api + + expected = ( + "AttestationMismatchError", + "DataFileAttestation", + "IndexAttestation", + "attest_data_file", + "attest_index", + "data_manifest_identity", + "file_generation_token", + "validate_data_file_attestation", + "validate_index_attestation", + ) + for name in expected: + assert getattr(dataset_api, name) is getattr(identity, name) + assert name in dataset_api.__all__ + assert "read_plan_identity" not in dataset_api.__all__ diff --git a/python/tsfile/dataset/__init__.py b/python/tsfile/dataset/__init__.py index 15c20d540..bbdcb6866 100644 --- a/python/tsfile/dataset/__init__.py +++ b/python/tsfile/dataset/__init__.py @@ -19,7 +19,32 @@ """Dataset-style TsFile accessors.""" from .dataframe import TsFileDataFrame +from .index_identity import ( + AttestationMismatchError, + DataFileAttestation, + IndexAttestation, + attest_data_file, + attest_index, + data_manifest_identity, + file_generation_token, + validate_data_file_attestation, + validate_index_attestation, +) from .metadata import SeriesPath from .timeseries import AlignedTimeseries, Timeseries -__all__ = ["TsFileDataFrame", "Timeseries", "AlignedTimeseries", "SeriesPath"] +__all__ = [ + "AlignedTimeseries", + "AttestationMismatchError", + "DataFileAttestation", + "IndexAttestation", + "SeriesPath", + "Timeseries", + "TsFileDataFrame", + "attest_data_file", + "attest_index", + "data_manifest_identity", + "file_generation_token", + "validate_data_file_attestation", + "validate_index_attestation", +] diff --git a/python/tsfile/dataset/dataframe.py b/python/tsfile/dataset/dataframe.py index c81990777..8d4bc1c66 100644 --- a/python/tsfile/dataset/dataframe.py +++ b/python/tsfile/dataset/dataframe.py @@ -25,7 +25,7 @@ import os import sys from types import SimpleNamespace -from typing import Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union import numpy as np @@ -43,6 +43,9 @@ from .merge import build_aligned_matrix, merge_time_value_parts, merge_timestamp_parts from .timeseries import AlignedTimeseries, Timeseries +if TYPE_CHECKING: + from .index_identity import DataFileAttestation, IndexAttestation + DeviceKey = Tuple[str, tuple] SeriesRefKey = Tuple[int, int] SeriesRef = Tuple[object, int, int] @@ -811,16 +814,41 @@ def __init__( show_progress: bool = True, use_index: bool = False, trust_index: Optional[bool] = None, + index_attestation: Optional["IndexAttestation"] = None, + data_file_attestations: Optional[Sequence["DataFileAttestation"]] = None, ): if not isinstance(use_index, bool): raise TypeError("use_index must be a bool") self._trust_index = _resolve_trust_index(trust_index) - self._paths = _expand_paths(paths) + has_index_attestation = index_attestation is not None + has_data_attestations = data_file_attestations is not None + if has_index_attestation != has_data_attestations: + raise ValueError( + "index_attestation and data_file_attestations must be provided together" + ) + if has_index_attestation and not self._trust_index: + raise ValueError("dataset attestations require trust_index=True") + + self._index_attestation = index_attestation + self._data_file_attestations = tuple(data_file_attestations or ()) + if has_index_attestation: + from .index_identity import DataFileAttestation, IndexAttestation + + if not isinstance(index_attestation, IndexAttestation): + raise TypeError("index_attestation must be an IndexAttestation") + if any( + not isinstance(item, DataFileAttestation) + for item in self._data_file_attestations + ): + raise TypeError( + "data_file_attestations must contain DataFileAttestation values" + ) + self._paths = [item.canonical_path for item in self._data_file_attestations] + else: + self._paths = _expand_paths(paths) self._show_progress = show_progress - # A trusted index is meaningful only when the persistent index path is - # used. Make the safe, explicit fast-path convenient for callers by - # enabling it automatically instead of requiring two flags. - self._use_index = use_index or self._trust_index + # Trusted and attested indexes always use the persistent mmap path. + self._use_index = use_index or self._trust_index or has_index_attestation self._readers: Dict[str, object] = {} self._index = _DataFrameCatalog() self._is_view = False @@ -843,6 +871,8 @@ def _from_subset( obj._show_progress = parent._show_progress obj._use_index = parent._use_index obj._trust_index = parent._trust_index + obj._index_attestation = parent._index_attestation + obj._data_file_attestations = parent._data_file_attestations obj._readers = parent._readers subset_refs = list(series_refs) obj._index = SimpleNamespace( @@ -943,11 +973,18 @@ def _load_metadata(self): ) from .runtime import DatasetRuntime - index_path = index_path_for(self._paths) + index_path = ( + self._index_attestation.canonical_index_path + if self._index_attestation is not None + else index_path_for(self._paths) + ) if self._trust_index: if not os.path.isfile(index_path): + index_kind = ( + "Attested" if self._index_attestation is not None else "Trusted" + ) raise FileNotFoundError( - f"Trusted Dataset Index not found: {index_path}" + f"{index_kind} Dataset Index not found: {index_path}" ) elif not index_matches_paths(index_path, self._paths): lock_path = index_path + ".lock" @@ -968,7 +1005,16 @@ def _load_metadata(self): reader.close() self._readers.clear() - self._runtime = DatasetRuntime(index_path, trust_index=self._trust_index) + self._runtime = DatasetRuntime( + index_path, + trust_index=self._trust_index, + index_attestation=self._index_attestation, + data_file_attestations=( + self._data_file_attestations + if self._index_attestation is not None + else None + ), + ) self._runtime_lease = self._runtime.lease() self._index = self._runtime.catalog if len(self._index.series) == 0: diff --git a/python/tsfile/dataset/index_identity.py b/python/tsfile/dataset/index_identity.py new file mode 100644 index 000000000..dbb687ba3 --- /dev/null +++ b/python/tsfile/dataset/index_identity.py @@ -0,0 +1,254 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Content and generation identities for trusted dataset indexes. + +Expensive content hashing belongs to the parent attestation pass. Workers use +the validation helpers, which compare canonical paths and stat metadata only. +The encodings below are explicitly versioned and length-prefixed so they do not +depend on Python object identity, mapping order, locale, or pickle. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import os +import struct +from typing import BinaryIO, Iterable + + +_INDEX_SCHEMA_VERSION = 1 +_DATA_MANIFEST_DOMAIN = b"tsfile-data-manifest-v1\0" +_FILE_GENERATION_DOMAIN = b"tsfile-file-generation-v1\0" +_DIGEST_HEX_LENGTH = 64 +_INDEX_FINGERPRINT_HEX_LENGTH = 16 +_DEFAULT_CHUNK_SIZE = 8 * 1024 * 1024 + + +class AttestationMismatchError(RuntimeError): + """Raised before reading when trusted metadata no longer matches.""" + + +@dataclass(frozen=True) +class IndexAttestation: + schema_version: int + canonical_index_path: str + index_identity_sha256: str + index_size: int + st_dev: int + st_ino: int + st_size: int + st_mtime_ns: int + + +@dataclass(frozen=True) +class DataFileAttestation: + file_id: int + canonical_path: str + st_dev: int + st_ino: int + st_size: int + st_mtime_ns: int + existing_index_fingerprint: str | None = None + + +def _canonical_path(path: os.PathLike[str] | str) -> str: + return os.path.realpath(os.path.abspath(os.fspath(path))) + + +def _stat_tuple(stat_result: os.stat_result) -> tuple[int, int, int, int]: + return ( + stat_result.st_dev, + stat_result.st_ino, + stat_result.st_size, + stat_result.st_mtime_ns, + ) + + +def _digest_bytes(value: str) -> bytes: + if ( + len(value) != _DIGEST_HEX_LENGTH + or value != value.lower() + or any(character not in "0123456789abcdef" for character in value) + ): + raise ValueError("identity must be a lowercase SHA-256 hex digest") + return bytes.fromhex(value) + + +def _optional_index_fingerprint_bytes(value: str | None) -> bytes: + if value is None: + return b"" + if ( + not isinstance(value, str) + or len(value) != _INDEX_FINGERPRINT_HEX_LENGTH + or value != value.lower() + or any(character not in "0123456789abcdef" for character in value) + ): + raise ValueError( + "existing_index_fingerprint must be a 16-character lowercase hex value" + ) + return value.encode("ascii") + + +def _u64(value: int, field: str) -> bytes: + if not isinstance(value, int) or isinstance(value, bool) or not 0 <= value < 2**64: + raise ValueError(f"{field} must fit uint64") + return struct.pack(" bytes: + return _u64(len(value), "field length") + value + + +def _update_stream(digest: "hashlib._Hash", stream: BinaryIO, chunk_size: int) -> None: + while chunk := stream.read(chunk_size): + digest.update(chunk) + + +def attest_index( + path: os.PathLike[str] | str, *, chunk_size: int = _DEFAULT_CHUNK_SIZE +) -> IndexAttestation: + """Hash finalized index bytes once in the trusted parent attestation pass.""" + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + canonical_path = _canonical_path(path) + digest = hashlib.sha256() + with open(canonical_path, "rb") as stream: + before_hash = os.fstat(stream.fileno()) + _update_stream(digest, stream, chunk_size) + after_hash = os.fstat(stream.fileno()) + if _stat_tuple(before_hash) != _stat_tuple(after_hash): + raise AttestationMismatchError("index changed while hashing") + return IndexAttestation( + schema_version=_INDEX_SCHEMA_VERSION, + canonical_index_path=canonical_path, + index_identity_sha256=digest.hexdigest(), + index_size=after_hash.st_size, + st_dev=after_hash.st_dev, + st_ino=after_hash.st_ino, + st_size=after_hash.st_size, + st_mtime_ns=after_hash.st_mtime_ns, + ) + + +def validate_index_attestation( + path: os.PathLike[str] | str, attestation: IndexAttestation +) -> None: + """Validate with one metadata lookup; never scan or re-hash index bytes.""" + canonical_path = _canonical_path(path) + if attestation.schema_version != _INDEX_SCHEMA_VERSION: + raise AttestationMismatchError("unsupported index attestation schema") + _digest_bytes(attestation.index_identity_sha256) + if canonical_path != attestation.canonical_index_path: + raise AttestationMismatchError("index path does not match attestation") + stat_result = os.stat(canonical_path) + expected = ( + attestation.st_dev, + attestation.st_ino, + attestation.st_size, + attestation.st_mtime_ns, + ) + if ( + attestation.index_size != attestation.st_size + or _stat_tuple(stat_result) != expected + ): + raise AttestationMismatchError("index stat does not match attestation") + + +def attest_data_file( + file_id: int, + path: os.PathLike[str] | str, + *, + existing_index_fingerprint: str | None = None, +) -> DataFileAttestation: + canonical_path = _canonical_path(path) + _u64(file_id, "file_id") + _optional_index_fingerprint_bytes(existing_index_fingerprint) + with open(canonical_path, "rb") as stream: + stat_result = os.fstat(stream.fileno()) + return DataFileAttestation( + file_id=file_id, + canonical_path=canonical_path, + st_dev=stat_result.st_dev, + st_ino=stat_result.st_ino, + st_size=stat_result.st_size, + st_mtime_ns=stat_result.st_mtime_ns, + existing_index_fingerprint=existing_index_fingerprint, + ) + + +def validate_data_file_attestation( + path: os.PathLike[str] | str, attestation: DataFileAttestation +) -> None: + canonical_path = _canonical_path(path) + if canonical_path != attestation.canonical_path: + raise AttestationMismatchError("data-file path does not match attestation") + stat_result = os.stat(canonical_path) + expected = ( + attestation.st_dev, + attestation.st_ino, + attestation.st_size, + attestation.st_mtime_ns, + ) + if _stat_tuple(stat_result) != expected: + raise AttestationMismatchError("data-file stat does not match attestation") + + +def _canonical_data_file(attestation: DataFileAttestation) -> bytes: + encoded_fingerprint = _optional_index_fingerprint_bytes( + attestation.existing_index_fingerprint + ) + return b"".join( + ( + _field(attestation.canonical_path.encode("utf-8")), + _u64(attestation.st_dev, "st_dev"), + _u64(attestation.st_ino, "st_ino"), + _u64(attestation.st_size, "st_size"), + _u64(attestation.st_mtime_ns, "st_mtime_ns"), + _field(encoded_fingerprint), + ) + ) + + +def data_manifest_identity( + attestations: Iterable[DataFileAttestation], *, require_sorted: bool = False +) -> str: + files = tuple(attestations) + if require_sorted and any( + left.file_id >= right.file_id for left, right in zip(files, files[1:]) + ): + raise ValueError("data-file IDs must be strictly increasing") + digest = hashlib.sha256(_DATA_MANIFEST_DOMAIN) + digest.update(_u64(len(files), "file count")) + for attestation in files: + digest.update(_u64(attestation.file_id, "file_id")) + digest.update(_field(_canonical_data_file(attestation))) + return digest.hexdigest() + + +def file_generation_token( + index_identity_sha256: str, + data_manifest_identity_sha256: str, + attestation: DataFileAttestation, +) -> str: + digest = hashlib.sha256(_FILE_GENERATION_DOMAIN) + digest.update(_digest_bytes(index_identity_sha256)) + digest.update(_digest_bytes(data_manifest_identity_sha256)) + digest.update(_u64(attestation.file_id, "file_id")) + digest.update(_field(_canonical_data_file(attestation))) + return digest.hexdigest() diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py index 7d62f66c3..44560c3cc 100644 --- a/python/tsfile/dataset/runtime.py +++ b/python/tsfile/dataset/runtime.py @@ -47,6 +47,15 @@ MappedDatasetIndex, file_fingerprint, ) +from .index_identity import ( + AttestationMismatchError, + DataFileAttestation, + IndexAttestation, + data_manifest_identity, + file_generation_token, + validate_data_file_attestation, + validate_index_attestation, +) from .metadata import ( MODEL_TABLE, MODEL_TREE, @@ -59,6 +68,10 @@ _SERIES_DESCRIPTOR_CACHE_SIZE = 4096 +def _digest_u64(value: str) -> int: + return int.from_bytes(bytes.fromhex(value)[:8], byteorder="little") + + def _configured_non_negative_int(name, default): raw = os.environ.get(name) if raw is None: @@ -240,6 +253,7 @@ def __init__( expected_size: int, fingerprint: int, validate_generation: bool = True, + data_file_attestation: Optional[DataFileAttestation] = None, ): self.file_id = file_id self.path = path @@ -248,6 +262,12 @@ def __init__( self._validate_generation_enabled = validate_generation if validate_generation: self._validate_generation() + if data_file_attestation is not None: + if data_file_attestation.file_id != file_id: + raise AttestationMismatchError( + "data-file attestation has the wrong file_id" + ) + validate_data_file_attestation(path, data_file_attestation) self.reader = TsFileReaderPy(path) self.active_uses = 0 self._closed = False @@ -278,10 +298,12 @@ def __init__( index: MappedDatasetIndex, max_open_files: int, validate_generation: bool = True, + data_file_attestations: Optional[Mapping[int, DataFileAttestation]] = None, ): self._index = index self.max_open_files = max(1, int(max_open_files)) self._validate_generation = bool(validate_generation) + self._data_file_attestations = dict(data_file_attestations or {}) self._sessions: "OrderedDict[int, _ReaderSession]" = OrderedDict() self._condition = threading.Condition() self._closed = False @@ -294,6 +316,7 @@ def _new_session(self, file_id: int): record[2], record[3], validate_generation=self._validate_generation, + data_file_attestation=self._data_file_attestations.get(file_id), ) @contextlib.contextmanager @@ -378,9 +401,15 @@ def __init__( index: MappedDatasetIndex, validate_references: bool = True, max_entries: Optional[int] = None, + mapped_index_identity: Optional[int] = None, + file_generation_tokens: Optional[Mapping[int, str]] = None, ): self._index = index self._validate_references = bool(validate_references) + self._mapped_index_identity = ( + id(index) if mapped_index_identity is None else int(mapped_index_identity) + ) + self._file_generation_tokens = dict(file_generation_tokens or {}) maximum = _SERIES_DESCRIPTOR_CACHE_SIZE if max_entries is None else max_entries if int(maximum) < 0: raise ValueError("max_entries must be non-negative") @@ -398,7 +427,7 @@ def _locator_tuple(self, file_id, locator_id): if self._validate_references and device_span[1] != file_id: raise ValueError("series locator points at another TsFile") return ( - id(self._index), + self._mapped_index_identity, file_id, file_record[2], file_record[3], @@ -464,7 +493,16 @@ def _release(self, entry): @contextlib.contextmanager def acquire(self, file_id, locator_id, reader, time_owner=None): - key = (id(self._index), file_id, locator_id) + generation = self._file_generation_tokens.get(file_id) + if generation is None: + file_record = self._index.record(TSFILE_RECORD, file_id) + generation = ( + self._mapped_index_identity, + file_id, + file_record[2], + file_record[3], + ) + key = (generation, locator_id) entry = None owner_entry = None prepared_time_owner = time_owner @@ -585,8 +623,27 @@ def __init__( query_workers: Optional[int] = None, query_parallel_min_rows: Optional[int] = None, trust_index: bool = False, + index_attestation: Optional[IndexAttestation] = None, + data_file_attestations: Optional[Sequence[DataFileAttestation]] = None, ): self.trust_index = bool(trust_index) + has_index_attestation = index_attestation is not None + has_data_attestations = data_file_attestations is not None + if has_index_attestation != has_data_attestations: + raise ValueError( + "index_attestation and data_file_attestations must be provided together" + ) + if has_index_attestation and not self.trust_index: + raise ValueError("attestations require trust_index=True") + if has_index_attestation and not isinstance( + index_attestation, IndexAttestation + ): + raise TypeError("index_attestation must be an IndexAttestation") + self.index_attestation = index_attestation + self.data_file_attestations = tuple(data_file_attestations or ()) + self.data_file_attestations_by_id = {} + self.data_manifest_identity_sha256 = None + self.file_generation_tokens = {} maximum = ( int(os.environ.get("TSFILE_DATAFRAME_MAX_OPEN_FILES", "16")) if max_open_files is None @@ -627,7 +684,21 @@ def __init__( self._torn_down = False try: + if self.index_attestation is not None: + validate_index_attestation(self._path, self.index_attestation) self.index = MappedDatasetIndex(self._path, trust_index=self.trust_index) + if self.index_attestation is not None: + expected_index_identity = ( + self.index_attestation.st_dev, + self.index_attestation.st_ino, + self.index_attestation.st_size, + self.index_attestation.st_mtime_ns, + ) + if self.index.identity != expected_index_identity: + raise AttestationMismatchError( + "mapped index does not match its attestation" + ) + self._configure_attestations() self._query_executor = ( ThreadPoolExecutor( max_workers=self.query_workers, @@ -640,11 +711,18 @@ def __init__( self.index, self.max_open_files, validate_generation=not self.trust_index, + data_file_attestations=self.data_file_attestations_by_id, ) self.prepared = PreparedSeriesCache( self.index, validate_references=not self.trust_index, max_entries=prepared_max_entries, + mapped_index_identity=( + _digest_u64(self.index_attestation.index_identity_sha256) + if self.index_attestation is not None + else None + ), + file_generation_tokens=self.file_generation_tokens, ) self.catalog = MappedDataFrameCatalog(self) except BaseException: @@ -654,6 +732,60 @@ def __init__( pass raise + def _configure_attestations(self): + if self.index_attestation is None: + return + if any( + not isinstance(attestation, DataFileAttestation) + for attestation in self.data_file_attestations + ): + raise TypeError( + "data_file_attestations must contain DataFileAttestation values" + ) + expected_file_ids = tuple(range(self.index.count(TSFILE_RECORD))) + actual_file_ids = tuple( + attestation.file_id for attestation in self.data_file_attestations + ) + if actual_file_ids != expected_file_ids: + raise ValueError( + "data_file_attestations must cover file_id=0..N-1 in order" + ) + + for attestation in self.data_file_attestations: + record = self.index.record(TSFILE_RECORD, attestation.file_id) + indexed_path = os.path.realpath( + os.path.abspath(self.index.string(record[0])) + ) + if ( + indexed_path != attestation.canonical_path + or record[2] != attestation.st_size + ): + raise AttestationMismatchError( + "data-file attestation does not match the Dataset Index" + ) + expected_fingerprint = f"{record[3]:016x}" + if attestation.existing_index_fingerprint != expected_fingerprint: + raise AttestationMismatchError( + "data-file fingerprint does not match the Dataset Index" + ) + + self.data_file_attestations_by_id = { + attestation.file_id: attestation + for attestation in self.data_file_attestations + } + self.data_manifest_identity_sha256 = data_manifest_identity( + self.data_file_attestations, + require_sorted=True, + ) + self.file_generation_tokens = { + attestation.file_id: file_generation_token( + self.index_attestation.index_identity_sha256, + self.data_manifest_identity_sha256, + attestation, + ) + for attestation in self.data_file_attestations + } + def _close_resources(self): if self._resources_closed: return @@ -681,6 +813,12 @@ def fork_replacement(self): query_workers=self.query_workers, query_parallel_min_rows=self.query_parallel_min_rows, trust_index=self.trust_index, + index_attestation=self.index_attestation, + data_file_attestations=( + self.data_file_attestations + if self.index_attestation is not None + else None + ), ) def _assert_current_process(self):