Refactor serialization buffer management and support memory tracking. #5231

teo-tsirpanis · 2024-08-08T00:40:42Z

This PR adds support for tracking memory used by serialization buffers. The first step towards implementing this is the replacement of the Buffer class with the new SerializationBuffer class, which has the following characteristics:

Like Buffer, SerializationBuffer supports both owned and non-owned memory buffers. This is necessary to maintain the semantics of the tiledb_buffer_t C API. SerializationBuffer also allows swapping between owned and non-owned modes.
SerializationBuffer does not support the write function which incrementally adds data by growing the buffer. Instead, it supports only the assign operation, which replaces the whole existing data of the buffer. assign copies the given data to a buffer owned by SerializationBuffer, unless a special marker value is passed, which changes it to just store a span to the buffer.
- There is also the assign_null_terminated method that sets an owned buffer and adds a null terminator at the end. This is used by JSON serialization for some reason (is it actually needed?).
SerializationBuffer does not store an offset and does not expose a read function that copies to a user-provided buffer and advances the offset. Instead the class is implicitly convertible to a const span of bytes, which gives access to the entire buffer at once.
SerializationBuffers are typically immutable between assignments, but some cases require subsequent modification. To cover these, the owned_mutable_span function returns a mutable span to the buffer, but will throw if used in a non-owned buffer.
SerializationBuffer supports being used with polymorphic allocators, which enables integration with the memory tracking system. A new MemoryTracker was added on ContextResources to track serialization-related allocations.

BufferList was also updated to store SerializationBuffers. Creations of serialization buffers and buffer lists were updated to pass an allocator, and C API handles for both types were likewise updated.

I also took up on the opportunity and converted all deserialization APIs that accepted a Buffer to accept a span<const char>, because they did not need any other facilities of SerializationBuffer.

TYPE: NO_HISTORY

Fixes compile errors due to redefinition of the name `span`.

And port `array_serialize` as a proof of concept.

…onBuffer`.

Only `SerializationBuffer` now implicitly converts to span.

This cleans-up code structure and will allow easily flowing the allocator from the list.

Oriskany

…cept an optional allocator.

`tiledb_buffer_list_t` was also updated to pass the allocator from the context's serialization memory tracker.

tiledb/sm/buffer/buffer.h

teo-tsirpanis · 2024-08-08T00:43:01Z

tiledb/sm/buffer/buffer.h

+  bool is_owned() const {
+    return buffer_.data() == buffer_owner_.data();
+  }


On second thought I could make this private. It has no use cases outside of this class. It could be used to guard for calling owned_mutable_span, but you are supposed to call this function on buffers you created yourself, and you know whether they are owned.

And reminder for me to document it.

tiledb/sm/buffer/buffer.h

tiledb/sm/rest/rest_client_remote.cc

teo-tsirpanis · 2024-08-08T00:54:55Z

tiledb/sm/buffer/buffer.h

+    if (!is_owned())
+      throw BufferStatusException(
+          "Cannot get a mutable span of a non-owned buffer.");


I could just remove that or replace it with assert, which would always return an empty span for non-owned buffers.

tiledb/api/c_api/buffer/buffer_api.cc

tiledb/sm/buffer/test/unit_serialization_buffer.cc

KiterLuc · 2024-08-22T08:34:48Z

tiledb/sm/rest/curl.cc

-  }
-
-  return num_read;
+  return buffer_list->read_at_most(dest, max_nbytes);


I think we should catch exceptions here and return CURL_READFUNC_ABORT?

read_at_most does not throw. I can further make it noexcept.

tiledb/sm/buffer/buffer.h

KiterLuc · 2024-08-22T09:32:21Z

tiledb/sm/buffer/buffer_list.cc


    // Read from buffer
    const uint64_t bytes_in_src = src.size() - current_relative_offset_;
    const uint64_t bytes_from_src = std::min(bytes_in_src, bytes_left);
    // If the destination pointer is not null, then read into it
    // if it is null then we are just seeking
    if (dest_ptr != nullptr)
-      RETURN_NOT_OK(src.read(dest_ptr + dest_offset, bytes_from_src));
+      memcpy(


Does this get rid of checks that were previously performed by read to make sure we don't read OOB?

Yes. The declaration of bytes_from_src above ensures we don't read OOB.

tiledb/sm/buffer/buffer_list.h

KiterLuc · 2024-08-22T10:40:10Z

tiledb/sm/buffer/test/unit_serialization_buffer.cc

+TEST_CASE(
+    "SerializationBuffer: Test memory tracking",
+    "[serialization_buffer][memory_tracking]") {
+  tiledb::sm::MemoryTrackerManager memory_tracker_manager;


Can we test what happens to copied/moved buffers related to memory tracking?

I wrote some tests locally and noticed the allocator does not propagate when copying a collection. 🤔

Good catch! Something we need to figure out for sure.

tiledb/sm/buffer/buffer.h

tiledb/sm/array/array.h

KiterLuc · 2024-09-04T16:02:33Z

tiledb/sm/serialization/query.cc

@@ -2738,23 +2729,12 @@ Status query_est_result_size_serialize(
      case SerializationType::JSON: {
        ::capnp::JsonCodec json;
        kj::String capnp_json = json.encode(est_result_size_builder);
-        const auto json_len = capnp_json.size();


Did we loose the behavior of adding \0 here?

And I was sure I would miss a case. Fixed, thanks!

tiledb/sm/serialization/group.cc

KiterLuc · 2024-09-04T16:26:44Z

tiledb/sm/serialization/query.cc

    SerializationType serialize_type,
    bool clientside,
    CopyState* copy_state,
    Query* query,
    ThreadPool* compute_tp) {
  // Create an original, serialized copy of the 'query' that we will revert
  // to if we are unable to deserialize 'serialized_buffer'.
-  BufferList original_bufferlist;
+  BufferList original_bufferlist(
+      query->resources().serialization_memory_tracker()->get_resource(


This is a very convoluted way to get the memory tracker. We have to find a way to pass it in query_deserialize without getting it from another object.

getting it from another object

My understanding is that this is how allocators propagate; from "parent" to "child" objects. We backup the state of a query in a temporary buffer list, and for me it makes sense to use the query's allocator.

We could get the allocator from the context argument passed in the C API, but I'd be against using it for cases where there are alternatives.

The serialized query buffer is not a child of the query IMO. We should do the work to properly pass the context into query_deserialize. It will require us to save the context instead of the threadpool in RestClientRemote.

Updated, we were already using the context in tiledb_deserialize_query.

KiterLuc · 2024-09-04T16:27:31Z

tiledb/sm/query/query.h

@@ -918,6 +918,13 @@ class Query {
   */
  RestClient* rest_client() const;

+  /**
+   * Returns the ContextResources associated to this query.
+   */


We should remove this. It was added as a way to get the serialization memory tracker but we should be able to get it from another source.

As said above, query is for me the best candidate to get the allocator.

tiledb/sm/rest/rest_client_remote.cc

teo-tsirpanis

@KiterLuc's changes LGTM, thanks!

teo-tsirpanis added 29 commits August 7, 2024 14:14

Update most deserialization APIs to accept the buffer in a span.

3794d0d

Use implicit conversion operator to go from buffer handle to span.

e936a82

Fixes compile errors due to redefinition of the name `span`.

Port some missed deserialization APIs to span.

fa99e60

Port query deserialization APIs to span.

5b9fabf

Add serialization buffer class.

9a6ad71

And port `array_serialize` as a proof of concept.

Update tiledb_buffer_t and serialization APIs to use a `Serializati…

8749ed8

…onBuffer`.

Remove implicit conversion of tiledb_buffer_t to span.

25dd27d

Only `SerializationBuffer` now implicitly converts to span.

Port some serialization APIs that were missed.

86e9d31

Update BufferList to hold serialization buffers.

62a97b8

Port the query serialization APIs to use SerializationBuffer.

3588eba

Fix compile errors.

50c3b3b

Update tests.

6588bc1

Move buffer list unit tests to standalone.

394809b

Fix constructors of SerializationBuffer.

ed465e1

Fix BufferList::read.

a6f07cb

Unstatus BufferList::read*** and seek.

6c3d34e

Construct serialization buffers in place in query serialization.

03e5674

This cleans-up code structure and will allow easily flowing the allocator from the list.

Add serialization buffer memory type.

cf4b441

Update query_serialize to accept a reference to a BufferList.

e99a47d

Oriskany

Add a serialization memory tracker and update tiledb_buffer_t to ac…

4a92d2e

…cept an optional allocator.

Update BufferList to hold a pmr vector of serialization buffers.

c3ec85a

`tiledb_buffer_list_t` was also updated to pass the allocator from the context's serialization memory tracker.

Instrument buffers used by the REST client.

64cfa20

Fix compile errors.

7b7cd3a

clang-format

48903a8

Fix heap memory violations.

65a31e6

Fix destruction of memory tracker with outstanding allocations.

a0b08e9

Add unit tests for SerializationBuffer.

28d26cd

Deallocate the owned buffer before reassignment.

2eb04e1

Fix tests.

418b0e8

teo-tsirpanis commented Aug 8, 2024

View reviewed changes

teo-tsirpanis requested a review from KiterLuc August 13, 2024 13:08

KiterLuc reviewed Aug 23, 2024

View reviewed changes

teo-tsirpanis added 4 commits August 23, 2024 14:16

Merge branch 'dev' into teo/serialization-buffer

6848cab

Update comments.

a55b2b6

Make allocator arguments mandatory.

8810c3f

Make more allocator arguments required and update usages.

1d3b67c

teo-tsirpanis force-pushed the teo/serialization-buffer branch from 9c8dcdf to 5162e5c Compare August 23, 2024 14:22

Fix more tests.

dad6d08

teo-tsirpanis force-pushed the teo/serialization-buffer branch from 5162e5c to dad6d08 Compare August 23, 2024 15:21

Fix segfault in C API buffer list test.

39bca77

KiterLuc reviewed Sep 4, 2024

View reviewed changes

tiledb/sm/array/array.h Outdated Show resolved Hide resolved

KiterLuc reviewed Sep 4, 2024

View reviewed changes

tiledb/sm/serialization/group.cc Show resolved Hide resolved

KiterLuc reviewed Sep 4, 2024

View reviewed changes

tiledb/sm/rest/rest_client_remote.cc Outdated Show resolved Hide resolved

teo-tsirpanis added 3 commits September 4, 2024 20:00

Merge branch 'dev' into teo/serialization-buffer

82fc192

Address most PR feedback.

e8fb876

Update query_deserialize to accept a memory tracker.

300298a

teo-tsirpanis force-pushed the teo/serialization-buffer branch from 8801934 to 300298a Compare September 4, 2024 19:28

Remove context resources accessors.

657e081

KiterLuc force-pushed the teo/serialization-buffer branch from 09eedef to 9d68068 Compare September 5, 2024 13:23

Remove moves and extra code.

a3fbc9a

KiterLuc force-pushed the teo/serialization-buffer branch from 9d68068 to a3fbc9a Compare September 5, 2024 15:29

teo-tsirpanis commented Sep 5, 2024

View reviewed changes

small comment fix.

21cf5d7

KiterLuc approved these changes Sep 6, 2024

View reviewed changes

KiterLuc merged commit ef7aba6 into dev Sep 6, 2024
62 checks passed

KiterLuc deleted the teo/serialization-buffer branch September 6, 2024 10:18

teo-tsirpanis mentioned this pull request Sep 12, 2024

The nightly GitHub Actions build job failed on Tuesday (2024-09-10) #5295

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor serialization buffer management and support memory tracking. #5231

Refactor serialization buffer management and support memory tracking. #5231

teo-tsirpanis commented Aug 8, 2024 •

edited by KiterLuc

Loading

teo-tsirpanis Aug 8, 2024

teo-tsirpanis Aug 8, 2024

KiterLuc Aug 22, 2024

teo-tsirpanis Aug 23, 2024

KiterLuc Aug 22, 2024

teo-tsirpanis Aug 23, 2024

KiterLuc Aug 22, 2024

teo-tsirpanis Aug 23, 2024 •

edited

Loading

KiterLuc Aug 26, 2024

KiterLuc Sep 4, 2024

teo-tsirpanis Sep 4, 2024

KiterLuc Sep 4, 2024

teo-tsirpanis Sep 4, 2024 •

edited

Loading

KiterLuc Sep 4, 2024

teo-tsirpanis Sep 4, 2024

KiterLuc Sep 4, 2024

teo-tsirpanis Sep 4, 2024

teo-tsirpanis left a comment

Refactor serialization buffer management and support memory tracking. #5231

Refactor serialization buffer management and support memory tracking. #5231

Conversation

teo-tsirpanis commented Aug 8, 2024 • edited by KiterLuc Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

teo-tsirpanis Aug 23, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

teo-tsirpanis Sep 4, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

teo-tsirpanis left a comment

Choose a reason for hiding this comment

teo-tsirpanis commented Aug 8, 2024 •

edited by KiterLuc

Loading

teo-tsirpanis Aug 23, 2024 •

edited

Loading

teo-tsirpanis Sep 4, 2024 •

edited

Loading