Skip to content

Commit 8867820

Browse files
laramielcopybara-github
laramiel
authored and
copybara-github
committed
Add logged trace spans to kvs backed chunk cache.
PiperOrigin-RevId: 744160481 Change-Id: Ib400a2f6f8c2bfdd7b11dbecb74594604861950e
1 parent 4ed1785 commit 8867820

File tree

7 files changed

+61
-9
lines changed

7 files changed

+61
-9
lines changed

‎tensorstore/driver/BUILD‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,14 +311,15 @@ tensorstore_cc_library(
311311
"//tensorstore:batch",
312312
"//tensorstore:box",
313313
"//tensorstore:chunk_layout",
314+
"//tensorstore:context",
314315
"//tensorstore:index",
315316
"//tensorstore:index_interval",
317+
"//tensorstore:json_serialization_options",
316318
"//tensorstore:open_mode",
317319
"//tensorstore:open_options",
318320
"//tensorstore:rank",
319321
"//tensorstore:resize_options",
320322
"//tensorstore:schema",
321-
"//tensorstore:spec",
322323
"//tensorstore:staleness_bound",
323324
"//tensorstore:transaction",
324325
"//tensorstore/index_space:index_transform",
@@ -361,6 +362,7 @@ tensorstore_cc_library(
361362
"//tensorstore/util:str_cat",
362363
"//tensorstore/util/execution",
363364
"//tensorstore/util/garbage_collection",
365+
"@com_github_nlohmann_json//:json",
364366
"@com_google_absl//absl/log:absl_log",
365367
"@com_google_absl//absl/meta:type_traits",
366368
"@com_google_absl//absl/status",

‎tensorstore/driver/kvs_backed_chunk_driver.h‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,39 +21,63 @@
2121
/// shared by multiple independent arrays) and one key-value store entry per
2222
/// chunk.
2323

24+
#include <stddef.h>
25+
26+
#include <functional>
2427
#include <memory>
28+
#include <optional>
29+
#include <string>
2530
#include <string_view>
31+
#include <utility>
2632

2733
#include "absl/status/status.h"
34+
#include "absl/strings/cord.h"
35+
#include "absl/time/time.h"
36+
#include <nlohmann/json_fwd.hpp>
37+
#include "tensorstore/batch.h"
2838
#include "tensorstore/box.h"
39+
#include "tensorstore/chunk_layout.h"
40+
#include "tensorstore/context.h"
2941
#include "tensorstore/driver/chunk_cache_driver.h"
3042
#include "tensorstore/driver/driver.h"
43+
#include "tensorstore/driver/driver_spec.h"
3144
#include "tensorstore/driver/registry.h"
3245
#include "tensorstore/index.h"
3346
#include "tensorstore/index_space/index_transform.h"
3447
#include "tensorstore/internal/cache/aggregate_writeback_cache.h"
3548
#include "tensorstore/internal/cache/async_cache.h"
3649
#include "tensorstore/internal/cache/async_initialized_cache_mixin.h"
50+
#include "tensorstore/internal/cache/cache.h"
3751
#include "tensorstore/internal/cache/cache_pool_resource.h"
38-
#include "tensorstore/internal/cache/chunk_cache.h"
3952
#include "tensorstore/internal/cache/kvs_backed_cache.h"
4053
#include "tensorstore/internal/cache/kvs_backed_chunk_cache.h"
4154
#include "tensorstore/internal/chunk_grid_specification.h"
4255
#include "tensorstore/internal/context_binding.h"
4356
#include "tensorstore/internal/data_copy_concurrency_resource.h"
57+
#include "tensorstore/internal/estimate_heap_usage/estimate_heap_usage.h"
4458
#include "tensorstore/internal/estimate_heap_usage/std_vector.h"
4559
#include "tensorstore/internal/intrusive_ptr.h"
4660
#include "tensorstore/internal/json_binding/bindable.h"
4761
#include "tensorstore/internal/open_mode_spec.h"
4862
#include "tensorstore/internal/type_traits.h"
63+
#include "tensorstore/json_serialization_options.h"
64+
#include "tensorstore/kvstore/driver.h"
4965
#include "tensorstore/kvstore/kvstore.h"
66+
#include "tensorstore/kvstore/spec.h"
5067
#include "tensorstore/open_mode.h"
68+
#include "tensorstore/open_options.h"
69+
#include "tensorstore/rank.h"
5170
#include "tensorstore/serialization/absl_time.h" // IWYU pragma: keep
5271
#include "tensorstore/serialization/std_optional.h" // IWYU pragma: keep
53-
#include "tensorstore/spec.h"
72+
#include "tensorstore/staleness_bound.h"
73+
#include "tensorstore/transaction.h"
74+
#include "tensorstore/util/dimension_set.h"
75+
#include "tensorstore/util/executor.h"
5476
#include "tensorstore/util/future.h"
77+
#include "tensorstore/util/garbage_collection/garbage_collection.h"
5578
#include "tensorstore/util/garbage_collection/std_optional.h" // IWYU pragma: keep
5679
#include "tensorstore/util/result.h"
80+
#include "tensorstore/util/span.h"
5781

5882
namespace tensorstore {
5983
namespace internal_kvs_backed_chunk_driver {

‎tensorstore/internal/cache/BUILD‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,7 @@ tensorstore_cc_library(
128128
"//tensorstore/internal/metrics:metadata",
129129
"//tensorstore/kvstore",
130130
"//tensorstore/kvstore:generation",
131+
"//tensorstore/util:future",
131132
"//tensorstore/util:status",
132133
"//tensorstore/util/execution",
133134
"//tensorstore/util/execution:any_receiver",
@@ -213,11 +214,14 @@ tensorstore_cc_library(
213214
"//tensorstore:array",
214215
"//tensorstore:index",
215216
"//tensorstore/internal:memory",
217+
"//tensorstore/internal/log:verbose_flag",
218+
"//tensorstore/internal/tracing",
216219
"//tensorstore/util:result",
217220
"//tensorstore/util:span",
218221
"//tensorstore/util:status",
219222
"//tensorstore/util:str_cat",
220223
"//tensorstore/util/execution",
224+
"@com_google_absl//absl/base:core_headers",
221225
"@com_google_absl//absl/container:fixed_array",
222226
"@com_google_absl//absl/container:inlined_vector",
223227
"@com_google_absl//absl/strings:cord",

‎tensorstore/internal/cache/kvs_backed_cache.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
#include "tensorstore/util/execution/any_receiver.h"
4444
#include "tensorstore/util/execution/execution.h"
4545
#include "tensorstore/util/execution/future_sender.h" // IWYU pragma: keep
46+
#include "tensorstore/util/future.h"
4647
#include "tensorstore/util/status.h"
4748

4849
namespace tensorstore {

‎tensorstore/internal/cache/kvs_backed_cache_testutil.cc‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,9 @@ void KvsBackedTestCache::Entry::DoDecode(std::optional<absl::Cord> value,
7171
return execution::set_error(
7272
receiver, absl::FailedPreconditionError("existing value contains Z"));
7373
}
74-
execution::set_value(
75-
receiver, std::make_shared<absl::Cord>(value.value_or(absl::Cord())));
74+
auto value_ptr = value ? std::make_shared<absl::Cord>(*std::move(value))
75+
: std::make_shared<absl::Cord>();
76+
return execution::set_value(receiver, value_ptr);
7677
}
7778

7879
void KvsBackedTestCache::Entry::DoEncode(std::shared_ptr<const absl::Cord> data,

‎tensorstore/internal/cache/kvs_backed_chunk_cache.cc‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#include <string>
2424
#include <utility>
2525

26+
#include "absl/base/attributes.h"
2627
#include "absl/container/fixed_array.h"
2728
#include "absl/container/inlined_vector.h"
2829
#include "absl/strings/cord.h"
@@ -32,7 +33,9 @@
3233
#include "tensorstore/internal/cache/cache.h"
3334
#include "tensorstore/internal/cache/chunk_cache.h"
3435
#include "tensorstore/internal/cache/kvs_backed_cache.h"
36+
#include "tensorstore/internal/log/verbose_flag.h"
3537
#include "tensorstore/internal/memory.h"
38+
#include "tensorstore/internal/tracing/logged_trace_span.h"
3639
#include "tensorstore/util/execution/execution.h"
3740
#include "tensorstore/util/result.h"
3841
#include "tensorstore/util/span.h"
@@ -41,6 +44,10 @@
4144

4245
namespace tensorstore {
4346
namespace internal {
47+
namespace {
48+
ABSL_CONST_INIT internal_log::VerboseFlag verbose_logging(
49+
"kvs_backed_chunk_cache");
50+
}
4451

4552
std::string KvsBackedChunkCache::Entry::GetKeyValueStoreKey() {
4653
auto& cache = GetOwningCache(*this);
@@ -56,12 +63,16 @@ void KvsBackedChunkCache::Entry::DoDecode(std::optional<absl::Cord> value,
5663
return;
5764
}
5865
auto& cache = GetOwningCache(*this);
66+
internal_tracing::LoggedTraceSpan trace_span(
67+
__func__, verbose_logging.Level(2),
68+
{{"cache", static_cast<void*>(&cache)}});
5969
auto decoded_result =
6070
cache.DecodeChunk(this->cell_indices(), *std::move(value));
6171
if (!decoded_result.ok()) {
62-
execution::set_error(receiver,
63-
internal::ConvertInvalidArgumentToFailedPrecondition(
64-
std::move(decoded_result).status()));
72+
auto status = internal::ConvertInvalidArgumentToFailedPrecondition(
73+
std::move(decoded_result).status());
74+
execution::set_error(
75+
receiver, std::move(trace_span).EndWithStatus(std::move(status)));
6576
return;
6677
}
6778
const size_t num_components = this->component_specs().size();
@@ -80,8 +91,12 @@ void KvsBackedChunkCache::Entry::DoEncode(std::shared_ptr<const ReadData> data,
8091
execution::set_value(receiver, std::nullopt);
8192
return;
8293
}
94+
8395
auto& entry = GetOwningEntry(*this);
8496
auto& cache = GetOwningCache(entry);
97+
internal_tracing::LoggedTraceSpan trace_span(
98+
__func__, verbose_logging.Level(2),
99+
{{"cache", static_cast<void*>(&cache)}});
85100
// Convert from array of `SharedArray<const void>` to array of
86101
// `SharedArrayView<const void>`.
87102
auto* components = data.get();
@@ -101,7 +116,9 @@ void KvsBackedChunkCache::Entry::DoEncode(std::shared_ptr<const ReadData> data,
101116
}
102117
auto encoded_result = cache.EncodeChunk(cell_indices, component_arrays);
103118
if (!encoded_result.ok()) {
104-
execution::set_error(receiver, std::move(encoded_result).status());
119+
execution::set_error(
120+
receiver, std::move(trace_span)
121+
.EndWithStatus(std::move(encoded_result).status()));
105122
return;
106123
}
107124
execution::set_value(receiver, *std::move(encoded_result));

‎tensorstore/internal/cache/kvs_backed_chunk_cache.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
#ifndef TENSORSTORE_INTERNAL_CACHE_KVS_BACKED_CHUNK_CACHE_H_
1616
#define TENSORSTORE_INTERNAL_CACHE_KVS_BACKED_CHUNK_CACHE_H_
1717

18+
#include <stddef.h>
19+
1820
#include <memory>
1921
#include <optional>
2022
#include <string>
@@ -23,6 +25,7 @@
2325
#include "absl/strings/cord.h"
2426
#include "tensorstore/array.h"
2527
#include "tensorstore/index.h"
28+
#include "tensorstore/internal/cache/async_cache.h"
2629
#include "tensorstore/internal/cache/chunk_cache.h"
2730
#include "tensorstore/internal/cache/kvs_backed_cache.h"
2831
#include "tensorstore/util/result.h"

0 commit comments

Comments
 (0)