Improve segment serialization for v1 segments (#2183)

#### Reference Issues/PRs
Ref Monday Ticket: 7852509418

#### What does this implement or fix?
This PR is fixing a segment fault that occurs in the
`CopyCompressedInterStoreTask` when used in enterprise.
The problem is a bit hard to reproduce as it involves:
1. Having some version V being written with an older version of ArcticDB
(e.g. 4.4.5) to some primary
2. Trying to replicate this version V to some secondary with a version
of enterprise that is using a newer version of ArcticDB (e.g. 5.2.3) and
that version should have a changed structure of the Segments (e.g. what
happened
[here](e3b5f53...757b4bd#diff-a17e4a0ad7760267af8c570795b362ec900136e414782a4b883485d36d9a8f84R167))
3. Because the newer version is trying to write more data than
allocated, we are getting a seg fault

The bug is that `CopyCompressedInterStoreTask` is trying to copy the
Segment but the information about its size is stale because it was
saved/reserialized based on the information from the older version.

The fix involves:
- Add null pointer checks in `decode_header_and_fields` and
`serialize_v1_header_to_buffer`
- the check in `decode_header_and_fields` is not needed for the fix but
it is a good change to prevent silly seg faults
- Simplify `calculate_size()` by removing unnecessary conditional
- Enhance buffer allocation and copying logic in
`serialize_v1_header_to_buffer`
- Add explicit size and pointer validation checks to prevent potential
memory issues

Technically, only [this
change](https://github.com/man-group/ArcticDB/pull/2183/files#diff-d35d87cb4a49c502a6519b27dc94590231057f633894d26c6c34f506fee1980fR154)
is needed for the fix of this bug.
But I think that the change in the
[segment.cpp](https://github.com/man-group/ArcticDB/pull/2183/files#diff-21f07ddd0b2b8ab0b82fd8392af726ab2a0b5ccafedab67d86497177d4f381a2R283)
make the code there easier to read/reason about + I've added some more
checks to make it more robust.

**Because the change is quite hard to test here, I've added the relevant
test for it in the enterprise repo with [this
PR](#2183

## Change Type (Required)
- [x] **Patch** (Bug fix or non-breaking improvement)
- [ ] **Minor** (New feature, but backward compatible)
- [ ] **Major** (Breaking changes)
- [ ] **Cherry pick**

#### Any other comments?

The fix is resolving seg faults like:
``` bash
==40129== Thread 123 IOPool0:
==40129== Invalid write of size 8
==40129==    at 0x4C2E8BE: memcpy@@GLIBC_2.14 (vg_replace_strmem.c:1035)
==40129==    by 0xC22DBCF1: arcticdb::Segment::serialize_v1_header_to_buffer(unsigned long) (segment.cpp:301)
==40129==    by 0xC22DBED6: arcticdb::Segment::serialize_header_v1() (segment.cpp:320)
==40129==    by 0xC22DBF5D: arcticdb::Segment::serialize_header() (segment.cpp:326)
==40129==    by 0xC3BD7A30: arcticdb::storage::s3::S3ClientImpl::put_object(std::string const&, arcticdb::Segment&, std::string const&, arcticdb::storage::s3::PutHeader) (s3_client_impl.cpp:192)
==40129==    by 0xC3BF1289: void arcticdb::storage::s3::detail::do_write_impl<arcticdb::storage::object_store_utils::FlatBucketizer>(arcticdb::storage::KeySegmentPair&, std::string const&, std::string const&, arcticdb::storage::s3::S3ClientInterface&, arcticdb::storage::object_store_utils::FlatBucketizer&&) (detail-inl.hpp:140)
==40129==    by 0xC3BE98A9: arcticdb::storage::s3::S3Storage::do_write(arcticdb::storage::KeySegmentPair&) (s3_storage.cpp:50)
==40129==    by 0xE608111: void arcticdb::storage::Storage::write<arcticdb::storage::KeySegmentPair&>(arcticdb::storage::KeySegmentPair&) (storage.hpp:48)
==40129==    by 0xE5F5C8B: arcticdb::storage::Storages::write(arcticdb::storage::KeySegmentPair&) (storages.hpp:46)
==40129==    by 0xE5F8245: arcticdb::storage::Library::write(arcticdb::storage::KeySegmentPair&) (library.hpp:90)
==40129==    by 0xE698D85: arcticdb::async::AsyncStore<arcticdb::util::SysClock>::write_compressed_sync(arcticdb::storage::KeySegmentPair) (async_store.hpp:180)
==40129==    by 0xE70ABDF: arcticdb::async::CopyCompressedInterStoreTask::copy() (tasks.hpp:402)
==40129==  Address 0x1416413b0 is 0 bytes after a block of size 1,151,856 alloc'd
==40129==    at 0x4C29F73: malloc (vg_replace_malloc.c:309)
==40129==    by 0xC373EE45: arcticdb::AllocatorImpl<arcticdb::NullTracingPolicy, arcticdb::util::LinearClock>::internal_alloc(unsigned long) (allocator.cpp:194)
==40129==    by 0xC373F143: arcticdb::AllocatorImpl<arcticdb::NullTracingPolicy, arcticdb::util::LinearClock>::aligned_alloc(unsigned long, bool) (allocator.cpp:303)
==40129==    by 0xC1DAD983: arcticdb::Buffer::resize(unsigned long) (buffer.hpp:232)
==40129==    by 0xC1DAD78A: arcticdb::Buffer::ensure(unsigned long) (buffer.hpp:177)
==40129==    by 0xC22DBBC5: arcticdb::Segment::serialize_v1_header_to_buffer(unsigned long) (segment.cpp:283)
==40129==    by 0xC22DBED6: arcticdb::Segment::serialize_header_v1() (segment.cpp:320)
==40129==    by 0xC22DBF5D: arcticdb::Segment::serialize_header() (segment.cpp:326)
==40129==    by 0xC3BD7A30: arcticdb::storage::s3::S3ClientImpl::put_object(std::string const&, arcticdb::Segment&, std::string const&, arcticdb::storage::s3::PutHeader) (s3_client_impl.cpp:192)
==40129==    by 0xC3BF1289: void arcticdb::storage::s3::detail::do_write_impl<arcticdb::storage::object_store_utils::FlatBucketizer>(arcticdb::storage::KeySegmentPair&, std::string const&, std::string const&, arcticdb::storage::s3::S3ClientInterface&, arcticdb::storage::object_store_utils::FlatBucketizer&&) (detail-inl.hpp:140)
==40129==    by 0xC3BE98A9: arcticdb::storage::s3::S3Storage::do_write(arcticdb::storage::KeySegmentPair&) (s3_storage.cpp:50)
==40129==    by 0xE608111: void arcticdb::storage::Storage::write<arcticdb::storage::KeySegmentPair&>(arcticdb::storage::KeySegmentPair&) (storage.hpp:48)
==40129== 
```

#### Checklist

<details>
  <summary>
   Checklist for code changes...
  </summary>
 
- [ ] Have you updated the relevant docstrings, documentation and
copyright notice?
- [ ] Is this contribution tested against [all ArcticDB's
features](../docs/mkdocs/docs/technical/contributing.md)?
- [ ] Do all exceptions introduced raise appropriate [error
messages](https://docs.arcticdb.io/error_messages/)?
 - [ ] Are API changes highlighted in the PR description?
- [ ] Is the PR labelled as enhancement or bug so it appears in
autogenerated release notes?
</details>

<!--
Thanks for contributing a Pull Request to ArcticDB! Please ensure you
have taken a look at:
- ArcticDB's Code of Conduct:
https://github.com/man-group/ArcticDB/blob/master/CODE_OF_CONDUCT.md
- ArcticDB's Contribution Licensing:
https://github.com/man-group/ArcticDB/blob/master/docs/mkdocs/docs/technical/contributing.md#contribution-licensing
-->

Author: G-D-Petrov

.gitignore

@@ -50,3 +50,6 @@ docs/mkdocs/docs/notebooks/.ipynb_checkpoints/
 # Ignore asv results
 python/.asv/html/
 python/.asv/results/
+
+# Ignore azurite files
+__azurite_*

cpp/arcticdb/codec/segment.cpp

@@ -139,6 +139,7 @@ struct DeserializedSegmentData {
 };
 
 DeserializedSegmentData decode_header_and_fields(const uint8_t*& src, bool copy_data) {
+    util::check(src != nullptr, "Got null data ptr from segment");
     auto* fixed_hdr = reinterpret_cast<const FixedHeader*>(src);
     ARCTICDB_DEBUG(log::codec(), "Reading header: {} + {} = {}", FIXED_HEADER_SIZE, fixed_hdr->header_bytes, FIXED_HEADER_SIZE + fixed_hdr->header_bytes);
 
@@ -279,14 +280,41 @@ std::pair<uint8_t*, size_t> Segment::serialize_v1_header_in_place(size_t total_h
 
 std::tuple<uint8_t*, size_t, std::unique_ptr<Buffer>> Segment::serialize_v1_header_to_buffer(size_t hdr_size) {
     auto tmp = std::make_unique<Buffer>();
-    ARCTICDB_TRACE(log::storage(), "Header doesn't fit in internal buffer, needed {} bytes but had {}, writing to temp buffer at {:x}", hdr_size, buffer_.preamble_bytes(),uintptr_t(tmp->data()));
-    tmp->ensure(calculate_size());
+    auto bytes_to_copy = buffer().bytes();
+    auto offset = FIXED_HEADER_SIZE + hdr_size;
+
+    auto total_size = offset + bytes_to_copy;
+
+    // Verify we have enough space for everything
+    tmp->ensure(total_size);
+    util::check(tmp->available() >= total_size,
+        "Buffer available space {} is less than required size {}", 
+        tmp->available(), 
+        total_size);
+
+    // This is both a sanity check and a way to populate the segment with the correct size
+    auto calculated_size = calculate_size();
+    util::check(total_size == calculated_size, "Expected total size {} to be equal to calculated size {}", total_size, calculated_size);
+
     auto* dst = tmp->preamble();
-    write_proto_header(dst);
-    std::memcpy(dst + FIXED_HEADER_SIZE + hdr_size,
-                buffer().data(),
-                buffer().bytes());
-    return std::make_tuple(tmp->preamble(), calculate_size(), std::move(tmp));
+    util::check(dst != nullptr, "Expected dst to be non-null");
+
+    auto written_hdr_size = write_proto_header(dst);
+    util::check(written_hdr_size == hdr_size, "Expected written header size {} to be equal to expected header size {}", written_hdr_size, hdr_size);
+
+    auto *final_dst = dst + offset;
+
+    auto *src = buffer().data();
+    if (src != nullptr) {   
+        std::memcpy(final_dst,
+                    src,
+                    bytes_to_copy);
+    } else {
+        util::check(bytes_to_copy == 0, "Expected bytes_to_copy to be 0 when src is nullptr");
+        ARCTICDB_DEBUG(log::codec(), "src is nullptr, skipping memcpy");
+    }
+
+    return std::make_tuple(dst, total_size, std::move(tmp));
 }
 
 std::tuple<uint8_t*, size_t, std::unique_ptr<Buffer>> Segment::serialize_header_v1() {

cpp/arcticdb/codec/segment.hpp

@@ -151,8 +151,7 @@ class Segment {
     }
 
     [[nodiscard]] std::size_t calculate_size() {
-        if(!size_.has_value())
-            size_ = FIXED_HEADER_SIZE + segment_header_bytes_size() + buffer_bytes();
+        size_ = FIXED_HEADER_SIZE + segment_header_bytes_size() + buffer_bytes();
 
         return *size_;
     }

cpp/arcticdb/toolbox/storage_mover.hpp

@@ -507,14 +507,25 @@ class ARCTICDB_VISIBILITY_HIDDEN StorageMover {
             size_t copied = 0;
             for (size_t offset = start; offset < end; ++offset) {
                 if (VariantKey& key = vkeys[offset]; source_store_->key_exists(key).get() && !target_store_->key_exists(key).get()) {
+                    util::check(variant_key_type(key) != KeyType::UNDEFINED, "Key type is undefined");
                     keys_to_copy[copied++] = std::pair{std::move(key), [copied, &segments](storage::KeySegmentPair&& ks) {
                         segments[copied] = std::move(ks);
                         return segments[copied].variant_key();
                     }};
+                } else {
+                    log::storage().warn("Key {} not found in source or already exists in target", key);
                 }
             }
+            // check that there are no undefined keys due to failed key_exists calls
+            std::erase_if(keys_to_copy, [](const auto& key) { return variant_key_type(key.first) == KeyType::UNDEFINED; });
+            if (keys_to_copy.empty()) {
+                continue;
+            }
+
             total_copied += copied;
-            [[maybe_unused]] auto keys = source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}).get();
+            auto keys = source_store_->batch_read_compressed(std::move(keys_to_copy), BatchReadArgs{}).get();
+            std::erase_if(segments, [](const auto& segment) { return variant_key_type(segment.variant_key()) == KeyType::UNDEFINED; });
+            util::check(keys.size() == segments.size(), "Keys and segments size mismatch, maybe due to parallel deletes");
             write_futs.push_back(target_store_->batch_write_compressed(std::move(segments)));
         }
         folly::collect(write_futs).get();

python/tests/compat/conftest.py python/arcticdb/util/venv.py

@@ -1,18 +1,15 @@
-import pytest
-import subprocess
-import os
-import venv
-import tempfile
 import logging
+import os
 import shutil
-from typing import Union, Optional, Dict, List
-from ..util.mark import (
-    AZURE_TESTS_MARK,
-    MONGO_TESTS_MARK,
-    VENV_COMPAT_TESTS_MARK,
-    PANDAS_2_COMPAT_TESTS_MARK
-)
+import subprocess
+import tempfile
+import venv
+
+from typing import Dict, List, Optional, Union
+
 from packaging.version import Version
+from arcticdb_ext import set_config_int, unset_config_int
+from arcticdb.arctic import Arctic
 
 logger = logging.getLogger("Compatibility tests")
 
@@ -45,8 +42,9 @@ def run_shell_command(
             stdin=subprocess.DEVNULL,
         )
     if result.returncode != 0:
-        logger.warning(
-            f"Command failed, stdout: {str(result.stdout)}, stderr: {str(result.stderr)}"
+        logger.error(
+            f"Command '{command_string}' failed with return code {result.returncode}\n"
+            f"stdout:\n{result.stdout.decode('utf-8')}\nstderr:\n{result.stderr.decode('utf-8')}"
         )
     return result
 
@@ -148,6 +146,10 @@ def create_library(self, lib_name: str) -> "VenvLib":
     def get_library(self, lib_name: str) -> "VenvLib":
         return VenvLib(self, lib_name, create_if_not_exists=False)
 
+    def cleanup(self):
+        ac = Arctic(self.uri)
+        for lib in ac.list_libraries():
+            ac.delete_library(lib)
 
 class VenvLib:
     def __init__(self, arctic, lib_name, create_if_not_exists=True):
@@ -177,68 +179,24 @@ def assert_read(self, sym: str, df) -> None:
         return self.execute(python_commands, {"expected_df": df})
 
 
-@pytest.fixture(
-    # scope="session",
-    params=[
-        pytest.param("1.6.2", marks=VENV_COMPAT_TESTS_MARK),
-        pytest.param("4.5.1", marks=VENV_COMPAT_TESTS_MARK),
-        pytest.param("5.0.0", marks=VENV_COMPAT_TESTS_MARK),
-    ],  # TODO: Extend this list with other old versions
-)
-def old_venv(request, tmp_path):
-    version = request.param
-    path = os.path.join("venvs", tmp_path, version)
-    compat_dir = os.path.dirname(os.path.abspath(__file__))
-    requirements_file = os.path.join(compat_dir, f"requirements-{version}.txt")
-    with Venv(path, requirements_file, version) as old_venv:
-        yield old_venv
-
-
-@pytest.fixture(
-    params=[
-        pytest.param("tmp_path", marks=PANDAS_2_COMPAT_TESTS_MARK)
-    ]
-)
-def pandas_v1_venv(request):
-    """A venv with Pandas v1 installed (and an old ArcticDB version). To help test compat across Pandas versions."""
-    version = "1.6.2"
-    tmp_path = request.getfixturevalue(request.param)
-    path = os.path.join("venvs", tmp_path, version)
-    compat_dir = os.path.dirname(os.path.abspath(__file__))
-    requirements_file = os.path.join(compat_dir, f"requirements-{version}.txt")
-    with Venv(path, requirements_file, version) as old_venv:
-        yield old_venv
-
-
-@pytest.fixture(
-    params=[
-        "lmdb",
-        "s3_ssl_disabled",
-        pytest.param("azurite", marks=AZURE_TESTS_MARK),
-        pytest.param("mongo", marks=MONGO_TESTS_MARK),
-    ]
-)
-def arctic_uri(request):
+class CurrentVersion:
     """
-    arctic_uri is a fixture which provides uri to all storages to be used for creating both old and current Arctic instances.
+    For many of the compatibility tests we need to maintain a single open connection to the library.
+    For example LMDB on Windows starts to fail if we at the same time we use an old_venv and current connection.
 
-    We use s3_ssl_disabled because which allows running tests for older versions like 1.6.2
+    So we use `with CurrentVersion` construct to ensure we delete all our outstanding references to the library.
     """
-    storage_fixture = request.getfixturevalue(request.param + "_storage")
-    if request.param == "mongo":
-        return storage_fixture.mongo_uri
-    else:
-        return storage_fixture.arctic_uri
-
-
-@pytest.fixture
-def old_venv_and_arctic_uri(old_venv, arctic_uri):
-    if arctic_uri.startswith("mongo") and "1.6.2" in old_venv.version:
-        pytest.skip("Mongo storage backend is not supported in 1.6.2")
+    def __init__(self, uri, lib_name):
+        self.uri = uri
+        self.lib_name = lib_name
 
-    if arctic_uri.startswith("lmdb") and Version(old_venv.version) < Version("5.0.0"):
-        pytest.skip(
-            "LMDB storage backed has a bug in versions before 5.0.0 which leads to flaky segfaults"
-        )
+    def __enter__(self):
+        set_config_int("VersionMap.ReloadInterval", 0) # We disable the cache to be able to read the data written from old_venv
+        self.ac = Arctic(self.uri)
+        self.lib = self.ac.get_library(self.lib_name)
+        return self
 
-    return old_venv, arctic_uri
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        unset_config_int("VersionMap.ReloadInterval")
+        del self.lib
+        del self.ac

python/tests/compat/arcticdb/test_compatibility.py

@@ -1,41 +1,17 @@
 import pytest
 from packaging import version
 import pandas as pd
-from arcticdb.arctic import Arctic
 from arcticdb.util.test import assert_frame_equal
-from arcticdb_ext import set_config_int, unset_config_int
 from arcticdb.options import ModifiableEnterpriseLibraryOption
 from arcticdb.toolbox.library_tool import LibraryTool
 from tests.util.mark import ARCTICDB_USING_CONDA
 
+from arcticdb.util.venv import CurrentVersion
 
 if ARCTICDB_USING_CONDA:
     pytest.skip("These tests rely on pip based environments", allow_module_level=True)
 
 
-class CurrentVersion:
-    """
-    For many of the compatibility tests we need to maintain a single open connection to the library.
-    For example LMDB on Windows starts to fail if we at the same time we use an old_venv and current connection.
-
-    So we use `with CurrentVersion` construct to ensure we delete all our outstanding references to the library.
-    """
-    def __init__(self, uri, lib_name):
-        self.uri = uri
-        self.lib_name = lib_name
-
-    def __enter__(self):
-        set_config_int("VersionMap.ReloadInterval", 0) # We disable the cache to be able to read the data written from old_venv
-        self.ac = Arctic(self.uri)
-        self.lib = self.ac.get_library(self.lib_name)
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        unset_config_int("VersionMap.ReloadInterval")
-        del self.lib
-        del self.ac
-
-
 def test_compat_write_read(old_venv_and_arctic_uri, lib_name):
     old_venv, arctic_uri = old_venv_and_arctic_uri
     sym = "sym"

python/tests/conftest.py

@@ -53,8 +53,12 @@
     MONGO_TESTS_MARK,
     REAL_S3_TESTS_MARK,
     SSL_TEST_SUPPORTED,
+    VENV_COMPAT_TESTS_MARK,
+    PANDAS_2_COMPAT_TESTS_MARK
 )
 from arcticdb.storage_fixtures.utils import safer_rmtree
+from packaging.version import Version
+from arcticdb.util.venv import Venv
 
 
 # region =================================== Misc. Constants & Setup ====================================
@@ -1089,3 +1093,74 @@ def in_memory_version_store_tiny_segment(in_memory_store_factory):
 @pytest.fixture(params=["lmdb_version_store_tiny_segment", "in_memory_version_store_tiny_segment"])
 def lmdb_or_in_memory_version_store_tiny_segment(request):
     return request.getfixturevalue(request.param)
+
+
+@pytest.fixture(
+    # scope="session",
+    params=[
+        pytest.param("1.6.2", marks=VENV_COMPAT_TESTS_MARK),
+        pytest.param("4.5.1", marks=VENV_COMPAT_TESTS_MARK),
+        pytest.param("5.0.0", marks=VENV_COMPAT_TESTS_MARK),
+    ],  # TODO: Extend this list with other old versions
+)
+def old_venv(request, tmp_path):
+    version = request.param
+    path = os.path.join("venvs", tmp_path, version)
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    compat_dir = os.path.join(test_dir, "compat")
+    requirements_file = os.path.join(compat_dir, f"requirements-{version}.txt")
+    with Venv(path, requirements_file, version) as old_venv:
+        yield old_venv
+
+
+@pytest.fixture(
+    params=[
+        pytest.param("tmp_path", marks=PANDAS_2_COMPAT_TESTS_MARK)
+    ]
+)
+def pandas_v1_venv(request):
+    """A venv with Pandas v1 installed (and an old ArcticDB version). To help test compat across Pandas versions."""
+    version = "1.6.2"
+    tmp_path = request.getfixturevalue(request.param)
+    path = os.path.join("venvs", tmp_path, version)
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    compat_dir = os.path.join(test_dir, "compat")
+    requirements_file = os.path.join(compat_dir, f"requirements-{version}.txt")
+    with Venv(path, requirements_file, version) as old_venv:
+        yield old_venv
+
+
+@pytest.fixture(
+    params=[
+        "lmdb",
+        "s3_ssl_disabled",
+        pytest.param("azurite", marks=AZURE_TESTS_MARK),
+        pytest.param("mongo", marks=MONGO_TESTS_MARK),
+    ]
+)
+def arctic_uri(request):
+    """
+    arctic_uri is a fixture which provides uri to all storages to be used for creating both old and current Arctic instances.
+
+    We use s3_ssl_disabled because which allows running tests for older versions like 1.6.2
+    """
+    storage_fixture = request.getfixturevalue(request.param + "_storage")
+    if request.param == "mongo":
+        return storage_fixture.mongo_uri
+    else:
+        return storage_fixture.arctic_uri
+
+
+@pytest.fixture
+def old_venv_and_arctic_uri(old_venv, arctic_uri):
+    if arctic_uri.startswith("mongo") and "1.6.2" in old_venv.version:
+        pytest.skip("Mongo storage backend is not supported in 1.6.2")
+
+    if arctic_uri.startswith("lmdb") and Version(old_venv.version) < Version("5.0.0"):
+        pytest.skip(
+            "LMDB storage backed has a bug in versions before 5.0.0 which leads to flaky segfaults"
+        )
+
+    yield old_venv, arctic_uri
+
+    old_venv.create_arctic(arctic_uri).cleanup()