Add file_io_mode context flag

"file_io_mode" indicates the mode used to read files. It currently supports the
values "default", "memmap" (replaces "file_io_memmap") and "direct" (linux only, experimental).

When set to "direct", if tensorstore detects that the reads can be satisified using direct io, then it will convert the file to use O_DIRECT and issue those reads.  Note that reads have to be aligned to and multiples of (typically) 4096 bytes, and this is only supported on relatively recent linux implementations.

"file_io_memmap" has been removed in favor of "file_io_mode": "memmap".  Otherwise this behaves the same.

PiperOrigin-RevId: 786790774
Change-Id: Ibdb9ec9f8523a44d5d8ef647fa298db328b47851

Author: laramiel

python/tensorstore/kvstore.cc

@@ -334,7 +334,7 @@ ensure there is a :python:`'/'` separator, use :py:obj:`.__truediv__` instead.
       'context': {
         'file_io_concurrency': {},
         'file_io_locking': {},
-        'file_io_memmap': False,
+        'file_io_mode': {},
         'file_io_sync': True,
       },
       'driver': 'file',
@@ -345,7 +345,7 @@ ensure there is a :python:`'/'` separator, use :py:obj:`.__truediv__` instead.
       'context': {
         'file_io_concurrency': {},
         'file_io_locking': {},
-        'file_io_memmap': False,
+        'file_io_mode': {},
         'file_io_sync': True,
       },
       'driver': 'file',
@@ -374,7 +374,7 @@ Returns a key-value store with an additional path component joined to the path.
       'context': {
         'file_io_concurrency': {},
         'file_io_locking': {},
-        'file_io_memmap': False,
+        'file_io_mode': {},
         'file_io_sync': True,
       },
       'driver': 'file',
@@ -385,7 +385,7 @@ Returns a key-value store with an additional path component joined to the path.
       'context': {
         'file_io_concurrency': {},
         'file_io_locking': {},
-        'file_io_memmap': False,
+        'file_io_mode': {},
         'file_io_sync': True,
       },
       'driver': 'file',
@@ -1174,7 +1174,7 @@ Returns a string representation based on the  :json:schema:`JSON representation<
       'context': {
         'file_io_concurrency': {},
         'file_io_locking': {},
-        'file_io_memmap': False,
+        'file_io_mode': {},
         'file_io_sync': True,
       },
       'driver': 'file',
@@ -1196,7 +1196,7 @@ Returns a copy of the key-value store.
     'context': {
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
     },
     'driver': 'file',
@@ -1207,7 +1207,7 @@ Returns a copy of the key-value store.
     'context': {
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
     },
     'driver': 'file',
@@ -1418,7 +1418,7 @@ Converts to the :json:schema:`JSON representation<KvStore>`.
    'driver': 'file',
    'file_io_concurrency': 'file_io_concurrency',
    'file_io_locking': 'file_io_locking',
-   'file_io_memmap': 'file_io_memmap',
+   'file_io_mode': 'file_io_mode',
    'file_io_sync': 'file_io_sync',
    'path': 'tmp/dataset/abc/'}
 

tensorstore/driver/auto/driver_test.cc

@@ -128,7 +128,7 @@ TEST(AutoTest, UnbindContextSimple) {
                 {"cache_pool", {{"total_bytes_limit", 1000}}},
                 {"file_io_concurrency", ::nlohmann::json::object_t()},
                 {"file_io_locking", ::nlohmann::json::object_t()},
-                {"file_io_memmap", false},
+                {"file_io_mode", ::nlohmann::json::object_t()},
                 {"file_io_sync", true},
                 {"ocdbt_coordinator", ::nlohmann::json::object_t()},
             }}})));
@@ -165,7 +165,7 @@ TEST(AutoTest, UnbindContext1) {
                 {"cache_pool", {{"total_bytes_limit", 1000}}},
                 {"file_io_concurrency", ::nlohmann::json::object_t()},
                 {"file_io_locking", ::nlohmann::json::object_t()},
-                {"file_io_memmap", false},
+                {"file_io_mode", ::nlohmann::json::object_t()},
                 {"file_io_sync", true},
                 {"ocdbt_coordinator", ::nlohmann::json::object_t()},
             }}})));

tensorstore/driver/auto/index.rst

@@ -67,7 +67,7 @@ A :ref:`zarr3<driver/zarr3>` TensorStore can be detected from its path:
       'data_copy_concurrency': {},
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
     },
     'driver': 'zarr3',
@@ -93,7 +93,7 @@ A :ref:`zarr3<driver/zarr3>` TensorStore can be detected from its path:
       'data_copy_concurrency': {},
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
     },
     'driver': 'zarr3',
@@ -169,7 +169,7 @@ used in conjunction with format auto-detection:
       'data_copy_concurrency': {},
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
     },
     'driver': 'cast',
@@ -196,7 +196,7 @@ from the path to the OCDBT database.
       'data_copy_concurrency': {},
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
       'ocdbt_coordinator': {},
     },
@@ -233,7 +233,7 @@ from the path to the OCDBT database.
       'data_copy_concurrency': {},
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
       'ocdbt_coordinator': {},
     },
@@ -278,7 +278,7 @@ of the OCDBT database:
       'data_copy_concurrency': {},
       'file_io_concurrency': {},
       'file_io_locking': {},
-      'file_io_memmap': False,
+      'file_io_mode': {},
       'file_io_sync': True,
       'ocdbt_coordinator': {},
     },

tensorstore/internal/os/BUILD

@@ -15,6 +15,28 @@ _LINKOPTS = select({
     "//conditions:default": [],
 })
 
+tensorstore_cc_library(
+    name = "aligned_alloc",
+    srcs = ["aligned_alloc.cc"],
+    hdrs = ["aligned_alloc.h"],
+    deps = [
+        ":include_windows",
+        ":memory_region",
+        "@abseil-cpp//absl/log:absl_log",
+    ],
+)
+
+tensorstore_cc_test(
+    name = "aligned_alloc_test",
+    srcs = ["aligned_alloc_test.cc"],
+    deps = [
+        ":aligned_alloc",
+        ":memory_region",
+        "@abseil-cpp//absl/strings:cord",
+        "@googletest//:gtest_main",
+    ],
+)
+
 tensorstore_cc_library(
     name = "cwd",
     srcs = ["cwd.cc"],
@@ -161,10 +183,12 @@ tensorstore_cc_test(
     name = "file_util_test",
     srcs = ["file_util_test.cc"],
     deps = [
+        ":aligned_alloc",
         ":file_util",
         "//tensorstore/internal/testing:scoped_directory",
         "//tensorstore/util:span",
         "//tensorstore/util:status_testutil",
+        "@abseil-cpp//absl/log:absl_log",
         "@abseil-cpp//absl/status",
         "@abseil-cpp//absl/strings",
         "@abseil-cpp//absl/strings:cord",

tensorstore/internal/os/aligned_alloc.cc

@@ -0,0 +1,71 @@
+// Copyright 2020 The TensorStore Authors
+//
+// Licensed 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.
+
+#include "tensorstore/internal/os/aligned_alloc.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cstdlib>
+
+#include "absl/log/absl_log.h"
+#include "tensorstore/internal/os/memory_region.h"
+
+// Include system headers last to reduce impact of macros.
+#include "tensorstore/internal/os/include_windows.h"
+
+namespace tensorstore {
+namespace internal_os {
+namespace {
+
+#ifdef _WIN32
+// https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/aligned-malloc?view=msvc-170
+void AlignedFree(char* data, size_t size) { _aligned_free(data); }
+#else
+void AlignedFree(char* data, size_t size) { free(data); }
+#endif
+
+}  // namespace
+
+MemoryRegion AllocatePageAlignedRegion(size_t alignment, size_t size) {
+  if (size == 0) {
+    return MemoryRegion(nullptr, 0, AlignedFree);
+  }
+  assert((alignment & (alignment - 1)) == 0);
+  size_t rounded_size = 0;
+
+#ifdef _WIN32
+  void* p = _aligned_malloc(size, alignment);
+#elif defined(__APPLE__)
+  // https://reviews.llvm.org/D138196
+  // ::aligned_alloc requires macos 15.0 or later. etc.
+  void* p = nullptr;
+  ::posix_memalign(&p, alignment, size);
+#else
+  rounded_size = (size + alignment - 1) & ~(alignment - 1);
+  void* p = ::aligned_alloc(alignment, std::max(size, rounded_size));
+#endif
+
+  if (p == nullptr) {
+    ABSL_LOG(FATAL) << "Failed to allocate memory " << size;
+  }
+  return MemoryRegion(static_cast<char*>(p), std::max(size, rounded_size),
+                      AlignedFree);
+}
+
+}  // namespace internal_os
+}  // namespace tensorstore

tensorstore/internal/os/aligned_alloc.h

@@ -0,0 +1,35 @@
+// Copyright 2020 The TensorStore Authors
+//
+// Licensed 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.
+
+#ifndef TENSORSTORE_INTERNAL_OS_ALIGNED_ALLOC_H_
+#define TENSORSTORE_INTERNAL_OS_ALIGNED_ALLOC_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <cstdio>
+#include <cstdlib>
+
+#include "tensorstore/internal/os/memory_region.h"
+
+namespace tensorstore {
+namespace internal_os {
+
+/// Try to allocate a region of memory backed the heap, page aligned.
+MemoryRegion AllocatePageAlignedRegion(size_t alignment, size_t size);
+
+}  // namespace internal_os
+}  // namespace tensorstore
+
+#endif  // TENSORSTORE_INTERNAL_OS_ALIGNED_ALLOC_H_

tensorstore/internal/os/aligned_alloc_test.cc

@@ -0,0 +1,41 @@
+// Copyright 2025 The TensorStore Authors
+//
+// Licensed 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.
+
+#include "tensorstore/internal/os/aligned_alloc.h"
+
+#include <stddef.h>
+
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/cord.h"
+#include "tensorstore/internal/os/memory_region.h"
+
+using ::tensorstore::internal_os::AllocatePageAlignedRegion;
+
+namespace {
+
+TEST(AlignedAllocTest, AllocatePageAlignedRegion) {
+  auto region = AllocatePageAlignedRegion(1024, 16 * 1024 * 1024);
+  EXPECT_THAT(region.as_string_view().size(), testing::Ge(16 * 1024 * 1024));
+
+  // Verify that assignment doesn't leak.
+  region = AllocatePageAlignedRegion(1024, 16);
+  EXPECT_THAT(region.as_string_view().size(), testing::Ge(16));
+
+  absl::Cord a = std::move(region).as_cord();
+}
+
+}  // namespace

tensorstore/internal/os/file_descriptor.h

@@ -17,27 +17,38 @@
 
 #include "tensorstore/internal/os/unique_handle.h"
 
+// Include system headers last to reduce impact of macros.
+#include "tensorstore/internal/os/include_windows.h"
+
 namespace tensorstore {
 namespace internal_os {
 
 #ifdef _WIN32
 // Specializations for Windows.
 
 /// Representation of open file/directory.
-using FileDescriptor = void*;  // HANDLE
+using FileDescriptor = HANDLE;  // HANDLE
+
+/// File descriptor traits for use with `UniqueHandle`.
+struct FileDescriptorTraits {
+  static FileDescriptor Invalid() { return INVALID_HANDLE_VALUE; }
+  static void Close(FileDescriptor fd);
+};
+
 #else
 // Specializations for Posix.
 
 /// Representation of open file/directory.
 using FileDescriptor = int;
-#endif
 
 /// File descriptor traits for use with `UniqueHandle`.
 struct FileDescriptorTraits {
-  static FileDescriptor Invalid() { return ((FileDescriptor)-1); }
+  static FileDescriptor Invalid() { return -1; }
   static void Close(FileDescriptor fd);
 };
 
+#endif
+
 /// Unique handle to an open file descriptor.
 ///
 /// The file descriptor is closed automatically by the destructor.

tensorstore/internal/os/file_util.h

@@ -67,6 +67,9 @@ class MappedRegion;
 /// multiple of this value.
 uint32_t GetDefaultPageSize();
 
+/// Returns the block alignment for Direct IO.
+size_t GetDirectIoBlockAlignment(FileDescriptor fd);
+
 /// Returns a Cord containing the contents of the file.
 ///
 /// \param fd File descriptor opened with `OpenFlags::OpenReadOnly`. The
@@ -87,6 +90,11 @@ enum class OpenFlags : int {
   Create = O_CREAT,
   Append = O_APPEND,
   Exclusive = O_EXCL,
+#if defined(O_DIRECT)
+  Direct = O_DIRECT,
+#else
+  Direct = 0x4000,
+#endif
   CloseOnExec = O_CLOEXEC,
   ReadWriteMask = O_RDONLY | O_WRONLY | O_RDWR,