Fix redundant guards in put_batch_internal and upsert_batch_internal bandwidth tracking

Agent-Logs-Url: https://github.com/kvcache-ai/Mooncake/sessions/51c607af-d065-40bd-b1b5-956c56cdc095

Co-authored-by: stmatengss <11641725+stmatengss@users.noreply.github.com>

mooncake-store/include/bandwidth_tracker.h

@@ -0,0 +1,51 @@
+#pragma once
+
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <string>
+#include <thread>
+
+namespace mooncake {
+
+// BandwidthTracker tracks read/write bandwidth for client operations.
+// Enabled via MC_ENABLE_BANDWIDTH_METRICS=1 environment variable.
+// Logging interval is configurable via MC_BANDWIDTH_LOG_INTERVAL_S (default 10
+// seconds).
+class BandwidthTracker {
+   public:
+    BandwidthTracker() = default;
+    ~BandwidthTracker();
+
+    // Start the bandwidth logging thread using the given client name.
+    // No-op if already started. Reads MC_ENABLE_BANDWIDTH_METRICS to decide
+    // whether to actually activate; returns immediately if disabled.
+    void start(const std::string& name);
+
+    // Stop the logging thread. No-op if not running.
+    void stop();
+
+    // Record bytes transferred in a write (put/upsert) operation.
+    // No-op when the tracker is not active.
+    void record_write(size_t bytes);
+
+    // Record bytes transferred in a read (get) operation.
+    // No-op when the tracker is not active.
+    void record_read(size_t bytes);
+
+   private:
+    void thread_func();
+
+    std::atomic<uint64_t> total_write_bytes_{0};
+    std::atomic<uint64_t> total_read_bytes_{0};
+    std::atomic<uint64_t> write_ops_{0};
+    std::atomic<uint64_t> read_ops_{0};
+
+    std::string name_;
+    int interval_s_{10};
+
+    std::thread thread_;
+    std::atomic<bool> running_{false};
+};
+
+}  // namespace mooncake

mooncake-store/include/dummy_client.h

@@ -4,6 +4,7 @@
 #include <csignal>
 #include <ylt/coro_rpc/coro_rpc_client.hpp>
 
+#include "bandwidth_tracker.h"
 #include "pyclient.h"
 #include "real_client.h"
 #include "shm_helper.h"
@@ -253,6 +254,9 @@ class DummyClient : public PyClient {
 
     // Ascend physical device id for dummy-real RPC to real, set in setup_dummy
     int32_t device_id_ = 0;
+
+    // Bandwidth metrics tracker (active when MC_ENABLE_BANDWIDTH_METRICS=1)
+    BandwidthTracker bandwidth_tracker_;
 };
 
 }  // namespace mooncake

mooncake-store/include/real_client.h

@@ -11,6 +11,7 @@
 #include <unordered_set>
 #include <vector>
 
+#include "bandwidth_tracker.h"
 #include "pyclient.h"
 #include "client_service.h"
 #include "client_buffer.hpp"
@@ -793,6 +794,9 @@ class RealClient : public PyClient {
     void teardown_ascend_shm_buffer(MappedShm &shm);
     tl::expected<void, ErrorCode> setup_ascend_internal(
         size_t local_buffer_size);
+
+    // Bandwidth metrics tracker (active when MC_ENABLE_BANDWIDTH_METRICS=1)
+    BandwidthTracker bandwidth_tracker_;
 };
 
 }  // namespace mooncake

mooncake-store/src/CMakeLists.txt

@@ -22,6 +22,7 @@ set(MOONCAKE_STORE_SOURCES
     aligned_client_buffer.cpp
     real_client.cpp
     dummy_client.cpp
+    bandwidth_tracker.cpp
     shm_helper.cpp
     http_metadata_server.cpp
     file_storage.cpp

mooncake-store/src/bandwidth_tracker.cpp

@@ -0,0 +1,148 @@
+#include "bandwidth_tracker.h"
+
+#include <cinttypes>
+#include <cstdio>
+#include <cstdlib>
+
+#include <glog/logging.h>
+
+#include "utils.h"
+
+namespace mooncake {
+
+namespace {
+
+std::string format_bw(double bytes_per_sec) {
+    char buf[64];
+    if (bytes_per_sec >= 1e9) {
+        snprintf(buf, sizeof(buf), "%.2f GB/s", bytes_per_sec / 1e9);
+    } else if (bytes_per_sec >= 1e6) {
+        snprintf(buf, sizeof(buf), "%.2f MB/s", bytes_per_sec / 1e6);
+    } else if (bytes_per_sec >= 1e3) {
+        snprintf(buf, sizeof(buf), "%.2f KB/s", bytes_per_sec / 1e3);
+    } else {
+        snprintf(buf, sizeof(buf), "%.2f B/s", bytes_per_sec);
+    }
+    return buf;
+}
+
+std::string format_bytes(uint64_t bytes) {
+    char buf[64];
+    if (bytes >= (1ULL << 30)) {
+        snprintf(buf, sizeof(buf), "%.2f GB",
+                 bytes / static_cast<double>(1ULL << 30));
+    } else if (bytes >= (1ULL << 20)) {
+        snprintf(buf, sizeof(buf), "%.2f MB",
+                 bytes / static_cast<double>(1ULL << 20));
+    } else if (bytes >= (1ULL << 10)) {
+        snprintf(buf, sizeof(buf), "%.2f KB",
+                 bytes / static_cast<double>(1ULL << 10));
+    } else {
+        snprintf(buf, sizeof(buf), "%" PRIu64 " B", bytes);
+    }
+    return buf;
+}
+
+}  // namespace
+
+BandwidthTracker::~BandwidthTracker() { stop(); }
+
+void BandwidthTracker::start(const std::string& name) {
+    // Check whether bandwidth metrics are enabled via environment variable.
+    if (!GetEnvOr<bool>("MC_ENABLE_BANDWIDTH_METRICS", false)) {
+        return;
+    }
+
+    // Prevent double-start.
+    if (running_.exchange(true)) {
+        return;
+    }
+
+    name_ = name;
+    interval_s_ = GetEnvOr<int>("MC_BANDWIDTH_LOG_INTERVAL_S", 10);
+
+    thread_ = std::thread([this] { thread_func(); });
+}
+
+void BandwidthTracker::stop() {
+    if (!running_.exchange(false)) {
+        return;
+    }
+    if (thread_.joinable()) {
+        thread_.join();
+    }
+}
+
+void BandwidthTracker::record_write(size_t bytes) {
+    if (!running_.load(std::memory_order_relaxed)) {
+        return;
+    }
+    total_write_bytes_.fetch_add(bytes, std::memory_order_relaxed);
+    write_ops_.fetch_add(1, std::memory_order_relaxed);
+}
+
+void BandwidthTracker::record_read(size_t bytes) {
+    if (!running_.load(std::memory_order_relaxed)) {
+        return;
+    }
+    total_read_bytes_.fetch_add(bytes, std::memory_order_relaxed);
+    read_ops_.fetch_add(1, std::memory_order_relaxed);
+}
+
+void BandwidthTracker::thread_func() {
+    using namespace std::chrono;
+
+    auto prev_time = steady_clock::now();
+    uint64_t prev_write = total_write_bytes_.load(std::memory_order_relaxed);
+    uint64_t prev_read = total_read_bytes_.load(std::memory_order_relaxed);
+    uint64_t prev_write_ops = write_ops_.load(std::memory_order_relaxed);
+    uint64_t prev_read_ops = read_ops_.load(std::memory_order_relaxed);
+
+    // Sleep in 100 ms chunks so stop() is responsive.
+    const int chunks = std::max(1, interval_s_ * 10);
+
+    while (running_.load(std::memory_order_relaxed)) {
+        for (int i = 0; i < chunks; ++i) {
+            if (!running_.load(std::memory_order_relaxed)) {
+                break;
+            }
+            std::this_thread::sleep_for(milliseconds(100));
+        }
+
+        if (!running_.load(std::memory_order_relaxed)) {
+            break;
+        }
+
+        auto now = steady_clock::now();
+        double elapsed =
+            duration_cast<duration<double>>(now - prev_time).count();
+        if (elapsed <= 0.0) {
+            continue;
+        }
+
+        uint64_t cur_write = total_write_bytes_.load(std::memory_order_relaxed);
+        uint64_t cur_read = total_read_bytes_.load(std::memory_order_relaxed);
+        uint64_t cur_write_ops = write_ops_.load(std::memory_order_relaxed);
+        uint64_t cur_read_ops = read_ops_.load(std::memory_order_relaxed);
+
+        double write_bw = (cur_write - prev_write) / elapsed;
+        double read_bw = (cur_read - prev_read) / elapsed;
+        double write_ops_rate = (cur_write_ops - prev_write_ops) / elapsed;
+        double read_ops_rate = (cur_read_ops - prev_read_ops) / elapsed;
+
+        LOG(INFO) << "[BandwidthMetrics][" << name_ << "] "
+                  << "Write: " << format_bw(write_bw) << " (" << write_ops_rate
+                  << " ops/s), " << "Read: " << format_bw(read_bw) << " ("
+                  << read_ops_rate << " ops/s) | "
+                  << "Total write: " << format_bytes(cur_write)
+                  << ", Total read: " << format_bytes(cur_read);
+
+        prev_time = now;
+        prev_write = cur_write;
+        prev_read = cur_read;
+        prev_write_ops = cur_write_ops;
+        prev_read_ops = cur_read_ops;
+    }
+}
+
+}  // namespace mooncake

mooncake-store/src/dummy_client.cpp

@@ -137,6 +137,8 @@ DummyClient::DummyClient() : client_id_(generate_uuid()) {
     client_pools_ =
         std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
             pool_conf);
+    // Start bandwidth metrics tracker if enabled via MC_ENABLE_BANDWIDTH_METRICS
+    bandwidth_tracker_.start("DummyClient");
 }
 
 DummyClient::~DummyClient() { tearDownAll(); }
@@ -524,35 +526,65 @@ uint64_t DummyClient::alloc_from_mem_pool(size_t size) {
 
 int DummyClient::put(const std::string& key, std::span<const char> value,
                      const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::put_dummy_helper, void>(
-        key, value, config, client_id_));
+    size_t write_bytes = value.size_bytes();
+    auto result = invoke_rpc<&RealClient::put_dummy_helper, void>(
+        key, value, config, client_id_);
+    if (result.has_value()) {
+        bandwidth_tracker_.record_write(write_bytes);
+    }
+    return to_py_ret(result);
 }
 
 int DummyClient::put_batch(const std::vector<std::string>& keys,
                            const std::vector<std::span<const char>>& values,
                            const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::put_batch_dummy_helper, void>(
-        keys, values, config, client_id_));
+    size_t write_bytes = 0;
+    for (const auto& v : values) {
+        write_bytes += v.size_bytes();
+    }
+    auto result = invoke_rpc<&RealClient::put_batch_dummy_helper, void>(
+        keys, values, config, client_id_);
+    if (result.has_value()) {
+        bandwidth_tracker_.record_write(write_bytes);
+    }
+    return to_py_ret(result);
 }
 
 int DummyClient::put_parts(const std::string& key,
                            std::vector<std::span<const char>> values,
                            const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::put_parts_dummy_helper, void>(
-        key, values, config, client_id_));
+    size_t write_bytes = 0;
+    for (const auto& v : values) {
+        write_bytes += v.size_bytes();
+    }
+    auto result = invoke_rpc<&RealClient::put_parts_dummy_helper, void>(
+        key, values, config, client_id_);
+    if (result.has_value()) {
+        bandwidth_tracker_.record_write(write_bytes);
+    }
+    return to_py_ret(result);
 }
 
 int DummyClient::upsert(const std::string& key, std::span<const char> value,
                         const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::upsert_dummy_helper, void>(
-        key, value, config, client_id_));
+    size_t write_bytes = value.size_bytes();
+    auto result = invoke_rpc<&RealClient::upsert_dummy_helper, void>(
+        key, value, config, client_id_);
+    if (result.has_value()) {
+        bandwidth_tracker_.record_write(write_bytes);
+    }
+    return to_py_ret(result);
 }
 
 int DummyClient::upsert_from(const std::string& key, void* buffer, size_t size,
                              const ReplicateConfig& config) {
     uint64_t dummy_addr = reinterpret_cast<uint64_t>(buffer);
-    return to_py_ret(invoke_rpc<&RealClient::upsert_from_dummy_helper, void>(
-        key, dummy_addr, size, config, client_id_));
+    auto result = invoke_rpc<&RealClient::upsert_from_dummy_helper, void>(
+        key, dummy_addr, size, config, client_id_);
+    if (result.has_value()) {
+        bandwidth_tracker_.record_write(size);
+    }
+    return to_py_ret(result);
 }
 
 std::vector<int> DummyClient::batch_upsert_from(
@@ -567,8 +599,11 @@ std::vector<int> DummyClient::batch_upsert_from(
             keys.size(), keys, buffers, sizes, config, client_id_);
     std::vector<int> results;
     results.reserve(internal_results.size());
-    for (const auto& result : internal_results) {
-        results.push_back(to_py_ret(result));
+    for (size_t i = 0; i < internal_results.size(); ++i) {
+        results.push_back(to_py_ret(internal_results[i]));
+        if (internal_results[i].has_value() && i < sizes.size()) {
+            bandwidth_tracker_.record_write(sizes[i]);
+        }
     }
     return results;
 }
@@ -576,15 +611,31 @@ std::vector<int> DummyClient::batch_upsert_from(
 int DummyClient::upsert_parts(const std::string& key,
                               std::vector<std::span<const char>> values,
                               const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::upsert_parts_dummy_helper, void>(
-        key, values, config, client_id_));
+    size_t write_bytes = 0;
+    for (const auto& v : values) {
+        write_bytes += v.size_bytes();
+    }
+    auto result = invoke_rpc<&RealClient::upsert_parts_dummy_helper, void>(
+        key, values, config, client_id_);
+    if (result.has_value()) {
+        bandwidth_tracker_.record_write(write_bytes);
+    }
+    return to_py_ret(result);
 }
 
 int DummyClient::upsert_batch(const std::vector<std::string>& keys,
                               const std::vector<std::span<const char>>& values,
                               const ReplicateConfig& config) {
-    return to_py_ret(invoke_rpc<&RealClient::upsert_batch_dummy_helper, void>(
-        keys, values, config, client_id_));
+    size_t write_bytes = 0;
+    for (const auto& v : values) {
+        write_bytes += v.size_bytes();
+    }
+    auto result = invoke_rpc<&RealClient::upsert_batch_dummy_helper, void>(
+        keys, values, config, client_id_);
+    if (result.has_value()) {
+        bandwidth_tracker_.record_write(write_bytes);
+    }
+    return to_py_ret(result);
 }
 
 int DummyClient::remove(const std::string& key, bool force) {
@@ -763,7 +814,11 @@ int64_t DummyClient::get_into(const std::string& key, void* buffer,
     if (!result) {
         return static_cast<int64_t>(toInt(result.error()));
     }
-    return to_py_ret(*result);
+    int64_t ret = to_py_ret(*result);
+    if (ret > 0) {
+        bandwidth_tracker_.record_read(static_cast<size_t>(ret));
+    }
+    return ret;
 }
 
 std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
@@ -787,6 +842,22 @@ std::vector<std::vector<std::vector<int64_t>>> DummyClient::get_into_ranges(
                                                internal_results.error());
     }
 
+    // Accumulate successfully read bytes for bandwidth metrics.
+    size_t read_bytes = 0;
+    for (const auto& buf_results : internal_results.value()) {
+        for (const auto& key_results : buf_results) {
+            for (const auto& fragment_result : key_results) {
+                if (fragment_result.has_value() && fragment_result.value() > 0) {
+                    read_bytes +=
+                        static_cast<size_t>(fragment_result.value());
+                }
+            }
+        }
+    }
+    if (read_bytes > 0) {
+        bandwidth_tracker_.record_read(read_bytes);
+    }
+
     return convert_ranged_read_results(internal_results.value());
 }
 
@@ -805,8 +876,11 @@ std::vector<int> DummyClient::batch_put_from(
     std::vector<int> results;
     results.reserve(internal_results.size());
 
-    for (const auto& result : internal_results) {
-        results.push_back(to_py_ret(result));
+    for (size_t i = 0; i < internal_results.size(); ++i) {
+        results.push_back(to_py_ret(internal_results[i]));
+        if (internal_results[i].has_value() && i < sizes.size()) {
+            bandwidth_tracker_.record_write(sizes[i]);
+        }
     }
 
     return results;
@@ -829,7 +903,11 @@ std::vector<int64_t> DummyClient::batch_get_into(
     results.reserve(internal_results.size());
 
     for (const auto& result : internal_results) {
-        results.push_back(to_py_ret(result));
+        int64_t ret = to_py_ret(result);
+        results.push_back(ret);
+        if (ret > 0) {
+            bandwidth_tracker_.record_read(static_cast<size_t>(ret));
+        }
     }
 
     return results;
@@ -856,8 +934,15 @@ std::vector<int> DummyClient::batch_put_from_multi_buffers(
                                config, device_id_, client_id_);
     std::vector<int> results;
     results.reserve(internal_results.size());
-    for (const auto& result : internal_results) {
-        results.push_back(to_py_ret(result));
+    for (size_t i = 0; i < internal_results.size(); ++i) {
+        results.push_back(to_py_ret(internal_results[i]));
+        if (internal_results[i].has_value() && i < all_sizes.size()) {
+            size_t key_bytes = 0;
+            for (const auto& s : all_sizes[i]) {
+                key_bytes += s;
+            }
+            bandwidth_tracker_.record_write(key_bytes);
+        }
     }
     return results;
 }
@@ -877,7 +962,11 @@ std::vector<int> DummyClient::batch_get_into_multi_buffers(
     std::vector<int> results;
     results.reserve(internal_results.size());
     for (const auto& result : internal_results) {
-        results.push_back(to_py_ret(result));
+        int64_t ret = to_py_ret(result);
+        results.push_back(static_cast<int>(ret));
+        if (ret > 0) {
+            bandwidth_tracker_.record_read(static_cast<size_t>(ret));
+        }
     }
     return results;
 }

mooncake-store/src/real_client.cpp

@@ -415,6 +415,8 @@ RealClient::RealClient() {
     mooncake::init_ylt_log_level();
     const char *hp = std::getenv("MC_STORE_USE_HUGEPAGE");
     use_hugepage_ = (hp != nullptr);
+    // Start bandwidth metrics tracker if enabled via MC_ENABLE_BANDWIDTH_METRICS
+    bandwidth_tracker_.start("RealClient");
 }
 
 RealClient::~RealClient() {
@@ -1015,6 +1017,7 @@ tl::expected<void, ErrorCode> RealClient::put_internal(
         return tl::unexpected(put_result.error());
     }
 
+    bandwidth_tracker_.record_write(value.size_bytes());
     return {};
 }
 
@@ -1097,11 +1100,14 @@ tl::expected<void, ErrorCode> RealClient::put_batch_internal(
     auto results = client_->BatchPut(keys, ordered_batched_slices, config);
 
     // Check if any operations failed
+    size_t total_bytes = 0;
     for (size_t i = 0; i < results.size(); ++i) {
         if (!results[i]) {
             return tl::unexpected(results[i].error());
         }
+        total_bytes += values[i].size_bytes();
     }
+    bandwidth_tracker_.record_write(total_bytes);
     return {};
 }
 
@@ -1185,6 +1191,7 @@ tl::expected<void, ErrorCode> RealClient::put_parts_internal(
         return tl::unexpected(put_result.error());
     }
 
+    bandwidth_tracker_.record_write(total_size);
     return {};
 }
 
@@ -2242,9 +2249,13 @@ tl::expected<int64_t, ErrorCode> RealClient::get_into_range_internal(
         return tl::unexpected(metadata_result.error());
     }
 
-    return execute_ranged_read(key, buffer, dst_offset, src_offset, size,
-                               metadata_result.value(),
-                               size_is_buffer_capacity);
+    auto result = execute_ranged_read(key, buffer, dst_offset, src_offset, size,
+                                      metadata_result.value(),
+                                      size_is_buffer_capacity);
+    if (result.has_value() && result.value() > 0) {
+        bandwidth_tracker_.record_read(static_cast<size_t>(result.value()));
+    }
+    return result;
 }
 
 int64_t RealClient::get_into(const std::string &key, void *buffer,
@@ -2479,7 +2490,13 @@ std::vector<tl::expected<void, ErrorCode>> RealClient::batch_put_from_internal(
     }
 
     // Call client BatchPut and return the vector<expected> directly
-    return client_->BatchPut(keys, ordered_batched_slices, config);
+    auto results = client_->BatchPut(keys, ordered_batched_slices, config);
+    for (size_t i = 0; i < results.size(); ++i) {
+        if (results[i].has_value() && i < sizes.size()) {
+            bandwidth_tracker_.record_write(sizes[i]);
+        }
+    }
+    return results;
 }
 
 tl::expected<void, ErrorCode> RealClient::put_from_internal(
@@ -2517,6 +2534,7 @@ tl::expected<void, ErrorCode> RealClient::put_from_internal(
         return tl::unexpected(put_result.error());
     }
 
+    bandwidth_tracker_.record_write(size);
     return {};
 }
 
@@ -2558,6 +2576,7 @@ tl::expected<void, ErrorCode> RealClient::upsert_internal(
     if (!result) {
         return tl::unexpected(result.error());
     }
+    bandwidth_tracker_.record_write(value.size_bytes());
     return {};
 }
 
@@ -2609,6 +2628,7 @@ tl::expected<void, ErrorCode> RealClient::upsert_from_internal(
     if (!result) {
         return tl::unexpected(result.error());
     }
+    bandwidth_tracker_.record_write(size);
     return {};
 }
 
@@ -2653,7 +2673,13 @@ RealClient::batch_upsert_from_internal(const std::vector<std::string> &keys,
         ordered_batched_slices.emplace_back(std::move(slices));
     }
 
-    return client_->BatchUpsert(keys, ordered_batched_slices, config);
+    auto results = client_->BatchUpsert(keys, ordered_batched_slices, config);
+    for (size_t i = 0; i < results.size(); ++i) {
+        if (results[i].has_value() && i < sizes.size()) {
+            bandwidth_tracker_.record_write(sizes[i]);
+        }
+    }
+    return results;
 }
 
 std::vector<int> RealClient::batch_upsert_from(
@@ -2778,6 +2804,7 @@ tl::expected<void, ErrorCode> RealClient::upsert_parts_internal(
                    << toString(result.error());
         return tl::unexpected(result.error());
     }
+    bandwidth_tracker_.record_write(total_size);
     return {};
 }
 
@@ -2861,11 +2888,14 @@ tl::expected<void, ErrorCode> RealClient::upsert_batch_internal(
     auto results = client_->BatchUpsert(keys, ordered_batched_slices, config);
 
     // Check if any operations failed
+    size_t total_bytes = 0;
     for (size_t i = 0; i < results.size(); ++i) {
         if (!results[i]) {
             return tl::unexpected(results[i].error());
         }
+        total_bytes += values[i].size_bytes();
     }
+    bandwidth_tracker_.record_write(total_bytes);
     return {};
 }
 
@@ -3277,6 +3307,16 @@ RealClient::batch_get_into_internal(const std::vector<std::string> &keys,
               << "us, with memory key count: " << valid_operations.size()
               << ", offload key count: " << offload_object_count;
 
+    size_t total_read_bytes = 0;
+    for (const auto &r : results) {
+        if (r.has_value() && r.value() > 0) {
+            total_read_bytes += static_cast<size_t>(r.value());
+        }
+    }
+    if (total_read_bytes > 0) {
+        bandwidth_tracker_.record_read(total_read_bytes);
+    }
+
     return results;
 }
 
@@ -3403,7 +3443,17 @@ RealClient::batch_put_from_multi_buffers_internal(
         }
     }
     // Call client BatchPut and return the vector<expected> directly
-    return client_->BatchPut(keys, batched_slices, config);
+    auto results = client_->BatchPut(keys, batched_slices, config);
+    for (size_t i = 0; i < results.size(); ++i) {
+        if (results[i].has_value() && i < all_sizes.size()) {
+            size_t key_bytes = 0;
+            for (const auto &s : all_sizes[i]) {
+                key_bytes += s;
+            }
+            bandwidth_tracker_.record_write(key_bytes);
+        }
+    }
+    return results;
 }
 
 std::vector<int> RealClient::batch_get_into_multi_buffers(
@@ -3556,6 +3606,17 @@ RealClient::batch_get_into_multi_buffers_internal(
             results[op.original_index] = tl::unexpected(error);
         }
     }
+
+    size_t total_read_bytes = 0;
+    for (const auto &r : results) {
+        if (r.has_value() && r.value() > 0) {
+            total_read_bytes += static_cast<size_t>(r.value());
+        }
+    }
+    if (total_read_bytes > 0) {
+        bandwidth_tracker_.record_read(total_read_bytes);
+    }
+
     return results;
 }