Support ExactKey queries in BlobDepot

5 files changed, 108 insertions, 73 deletions

diff --git a/ydb/core/blob_depot/agent/blob_mapping_cache.cpp b/ydb/core/blob_depot/agent/blob_mapping_cache.cpp
index cb20a439d58..66253324126 100644
--- a/ydb/core/blob_depot/agent/blob_mapping_cache.cpp
+++ b/ydb/core/blob_depot/agent/blob_mapping_cache.cpp
@@ -60,9 +60,7 @@ namespace NKikimr::NBlobDepot {
 
             NKikimrBlobDepot::TEvResolve msg;
             auto *item = msg.AddItems();
-            item->SetBeginningKey(it->first);
-            item->SetEndingKey(it->first);
-            item->SetIncludeEnding(true);
+            item->SetExactKey(it->first);
 
             if (Agent.VirtualGroupId) {
                 const auto& id = TLogoBlobID::FromBinary(it->first);
diff --git a/ydb/core/blob_depot/agent/storage_discover.cpp b/ydb/core/blob_depot/agent/storage_discover.cpp
index 2055a48ec76..e8f9b66d6d7 100644
--- a/ydb/core/blob_depot/agent/storage_discover.cpp
+++ b/ydb/core/blob_depot/agent/storage_discover.cpp
@@ -49,12 +49,13 @@ namespace NKikimr::NBlobDepot {
 
                 NKikimrBlobDepot::TEvResolve resolve;
                 auto *item = resolve.AddItems();
-                item->SetBeginningKey(from.AsBinaryString());
-                item->SetIncludeBeginning(true);
-                item->SetEndingKey(to.AsBinaryString());
-                item->SetIncludeEnding(true);
-                item->SetMaxKeys(1);
-                item->SetReverse(true);
+                auto *range = item->MutableKeyRange();
+                range->SetBeginningKey(from.AsBinaryString());
+                range->SetIncludeBeginning(true);
+                range->SetEndingKey(to.AsBinaryString());
+                range->SetIncludeEnding(true);
+                range->SetMaxKeys(1);
+                range->SetReverse(true);
                 item->SetTabletId(TabletId);
                 item->SetMustRestoreFirst(true);
 
diff --git a/ydb/core/blob_depot/agent/storage_range.cpp b/ydb/core/blob_depot/agent/storage_range.cpp
index 7e0143414fd..7feea1c19ef 100644
--- a/ydb/core/blob_depot/agent/storage_range.cpp
+++ b/ydb/core/blob_depot/agent/storage_range.cpp
@@ -21,7 +21,7 @@ namespace NKikimr::NBlobDepot {
             using TQuery::TQuery;
 
             void Initiate() override {
-                auto& msg = *Event->Get<TEvBlobStorage::TEvRange>();
+                auto& msg = GetQuery();
 
                 if (msg.Decommission) {
                     Y_VERIFY(Agent.ProxyId);
@@ -40,7 +40,7 @@ namespace NKikimr::NBlobDepot {
             }
 
             void IssueResolve() {
-                auto& msg = *Event->Get<TEvBlobStorage::TEvRange>();
+                auto& msg = GetQuery();
 
                 TString from = msg.From.AsBinaryString();
                 TString to = msg.To.AsBinaryString();
@@ -51,11 +51,12 @@ namespace NKikimr::NBlobDepot {
 
                 NKikimrBlobDepot::TEvResolve resolve;
                 auto *item = resolve.AddItems();
-                item->SetBeginningKey(from);
-                item->SetIncludeBeginning(true);
-                item->SetEndingKey(to);
-                item->SetIncludeEnding(true);
-                item->SetReverse(reverse);
+                auto *range = item->MutableKeyRange();
+                range->SetBeginningKey(from);
+                range->SetIncludeBeginning(true);
+                range->SetEndingKey(to);
+                range->SetIncludeEnding(true);
+                range->SetReverse(reverse);
                 item->SetTabletId(msg.TabletId);
                 item->SetMustRestoreFirst(msg.MustRestoreFirst);
 
@@ -64,12 +65,12 @@ namespace NKikimr::NBlobDepot {
             }
 
             void IssueResolve(TLogoBlobID id, size_t index) {
+                auto& msg = GetQuery();
                 NKikimrBlobDepot::TEvResolve resolve;
                 auto *item = resolve.AddItems();
-                item->SetBeginningKey(id.AsBinaryString());
-                item->SetIncludeBeginning(true);
-                item->SetEndingKey(id.AsBinaryString());
-                item->SetIncludeEnding(true);
+                item->SetExactKey(id.AsBinaryString());
+                item->SetTabletId(msg.TabletId);
+                item->SetMustRestoreFirst(msg.MustRestoreFirst);
 
                 Agent.Issue(std::move(resolve), this, std::make_shared<TExtraResolveContext>(index));
                 ++ResolvesInFlight;
@@ -88,7 +89,7 @@ namespace NKikimr::NBlobDepot {
             }
 
             void HandleResolveResult(ui64 id, TRequestContext::TPtr context, NKikimrBlobDepot::TEvResolveResult& msg) {
-                auto& query = *Event->Get<TEvBlobStorage::TEvRange>();
+                auto& query = GetQuery();
 
                 --ResolvesInFlight;
 
@@ -153,6 +154,10 @@ namespace NKikimr::NBlobDepot {
                     EndWithSuccess(std::move(Response));
                 }
             }
+
+            TEvBlobStorage::TEvRange& GetQuery() const {
+                return *Event->Get<TEvBlobStorage::TEvRange>();
+            }
         };
 
         return new TRangeQuery(*this, std::move(ev));
diff --git a/ydb/core/blob_depot/data_resolve.cpp b/ydb/core/blob_depot/data_resolve.cpp
index 4f9464d6d30..92629532f9a 100644
--- a/ydb/core/blob_depot/data_resolve.cpp
+++ b/ydb/core/blob_depot/data_resolve.cpp
@@ -29,6 +29,43 @@ namespace NKikimr::NBlobDepot {
             , NumKeysRead(predecessor.NumKeysRead)
         {}
 
+        bool GetScanParams(const NKikimrBlobDepot::TEvResolve::TItem& item, std::optional<TKey> *begin,
+                std::optional<TKey> *end, TScanFlags *flags, ui64 *maxKeys) {
+            switch (item.GetKeyDesignatorCase()) {
+                case NKikimrBlobDepot::TEvResolve::TItem::kKeyRange: {
+                    const auto& range = item.GetKeyRange();
+                    *flags = TScanFlags()
+                        | (range.GetIncludeBeginning() ? EScanFlags::INCLUDE_BEGIN : TScanFlags())
+                        | (range.GetIncludeEnding() ? EScanFlags::INCLUDE_END : TScanFlags())
+                        | (range.GetReverse() ? EScanFlags::REVERSE : TScanFlags());
+                    if (range.HasBeginningKey()) {
+                        begin->emplace(TKey::FromBinaryKey(range.GetBeginningKey(), Self->Config));
+                    } else {
+                        begin->reset();
+                    }
+                    if (range.HasEndingKey()) {
+                        end->emplace(TKey::FromBinaryKey(range.GetEndingKey(), Self->Config));
+                    } else {
+                        end->reset();
+                    }
+                    *maxKeys = range.GetMaxKeys();
+                    return true;
+                }
+
+                case NKikimrBlobDepot::TEvResolve::TItem::kExactKey:
+                    begin->emplace(TKey::FromBinaryKey(item.GetExactKey(), Self->Config));
+                    end->emplace(begin->value());
+                    *flags = EScanFlags::INCLUDE_BEGIN | EScanFlags::INCLUDE_END;
+                    *maxKeys = 1;
+                    return true;
+
+                case NKikimrBlobDepot::TEvResolve::TItem::KEYDESIGNATOR_NOT_SET:
+                    return false;
+            }
+
+            return false;
+        }
+
         bool Execute(TTransactionContext& txc, const TActorContext&) override {
             NIceDb::TNiceDb db(txc.DB);
 
@@ -47,24 +84,12 @@ namespace NKikimr::NBlobDepot {
             for (; ItemIndex < items.size(); ++ItemIndex, LastScannedKey.reset(), NumKeysRead = 0) {
                 const auto& item = items[ItemIndex];
 
-                std::optional<TKey> begin = item.HasBeginningKey()
-                    ? std::make_optional(TKey::FromBinaryKey(item.GetBeginningKey(), Self->Config))
-                    : std::nullopt;
-
-                std::optional<TKey> end = item.HasEndingKey()
-                    ? std::make_optional(TKey::FromBinaryKey(item.GetEndingKey(), Self->Config))
-                    : std::nullopt;
-
+                std::optional<TKey> begin;
+                std::optional<TKey> end;
                 TScanFlags flags;
-                if (item.GetIncludeBeginning()) {
-                    flags |= EScanFlags::INCLUDE_BEGIN;
-                }
-                if (item.GetIncludeEnding()) {
-                    flags |= EScanFlags::INCLUDE_END;
-                }
-                if (item.GetReverse()) {
-                    flags |= EScanFlags::REVERSE;
-                }
+                ui64 maxKeys;
+                const bool success = GetScanParams(item, &begin, &end, &flags, &maxKeys);
+                Y_VERIFY_DEBUG(success);
 
                 // adjust range according to actually generated data
                 if (LastScannedKey) {
@@ -84,7 +109,7 @@ namespace NKikimr::NBlobDepot {
                     // we can scan only some part from memory -- do it
                     auto callback = [&](const TKey& key, const TValue&) {
                         LastScannedKey = key;
-                        return ++NumKeysRead != item.GetMaxKeys();
+                        return ++NumKeysRead != maxKeys;
                     };
                     Self->Data->ScanRange(begin ? &begin.value() : nullptr, &Self->Data->LastLoadedKey.value(),
                         flags | EScanFlags::INCLUDE_END, callback);
@@ -94,7 +119,7 @@ namespace NKikimr::NBlobDepot {
                     flags &= ~EScanFlags::INCLUDE_BEGIN;
 
                     // check if we have read all the keys requested
-                    if (NumKeysRead == item.GetMaxKeys()) {
+                    if (NumKeysRead == maxKeys) {
                         continue;
                     }
                 }
@@ -115,7 +140,7 @@ namespace NKikimr::NBlobDepot {
 
                             const bool matchBegin = !begin || (flags & EScanFlags::INCLUDE_BEGIN ? *begin <= key : *begin < key);
                             const bool matchEnd = !end || (flags & EScanFlags::INCLUDE_END ? key <= *end : key < *end);
-                            if (matchBegin && matchEnd && ++NumKeysRead == item.GetMaxKeys()) {
+                            if (matchBegin && matchEnd && ++NumKeysRead == maxKeys) {
                                 // we have hit the MaxItems limit, exit
                                 return true;
                             } else if (flags & EScanFlags::REVERSE ? !matchBegin : !matchEnd) {
@@ -137,7 +162,7 @@ namespace NKikimr::NBlobDepot {
                         : applyEnd(std::forward<std::decay_t<decltype(x)>>(x));
                 };
                 auto applyReverse = [&](auto&& x) {
-                    return item.GetReverse()
+                    return flags & EScanFlags::REVERSE
                         ? applyBegin(x.Reverse())
                         : applyBegin(std::forward<std::decay_t<decltype(x)>>(x));
                 };
@@ -186,26 +211,12 @@ namespace NKikimr::NBlobDepot {
             for (const auto& item : Request->Get()->Record.GetItems()) {
                 std::optional<ui64> cookie = item.HasCookie() ? std::make_optional(item.GetCookie()) : std::nullopt;
 
-                std::optional<TKey> begin = item.HasBeginningKey()
-                    ? std::make_optional(TKey::FromBinaryKey(item.GetBeginningKey(), Self->Config))
-                    : std::nullopt;
-
-                std::optional<TKey> end = item.HasEndingKey()
-                    ? std::make_optional(TKey::FromBinaryKey(item.GetEndingKey(), Self->Config))
-                    : std::nullopt;
-
+                std::optional<TKey> begin;
+                std::optional<TKey> end;
                 TScanFlags flags;
-                if (item.GetIncludeBeginning()) {
-                    flags |= EScanFlags::INCLUDE_BEGIN;
-                }
-                if (item.GetIncludeEnding()) {
-                    flags |= EScanFlags::INCLUDE_END;
-                }
-                if (item.GetReverse()) {
-                    flags |= EScanFlags::REVERSE;
-                }
-
-                ui64 count = item.GetMaxKeys();
+                ui64 count;
+                const bool success = GetScanParams(item, &begin, &end, &flags, &count);
+                Y_VERIFY_DEBUG(success);
 
                 auto callback = [&](const TKey& key, const TValue& value) {
                     IssueResponseItem(cookie, key, value, lastResponseSize);
@@ -292,11 +303,25 @@ namespace NKikimr::NBlobDepot {
                 TLogoBlobID maxId(tabletId, Max<ui32>(), Max<ui32>(), TLogoBlobID::MaxChannel, TLogoBlobID::MaxBlobSize,
                     TLogoBlobID::MaxCookie, TLogoBlobID::MaxPartId, TLogoBlobID::MaxCrcMode);
 
-                if (item.HasBeginningKey()) {
-                    minId = TKey::FromBinaryKey(item.GetBeginningKey(), Self->Config).GetBlobId();
-                }
-                if (item.HasEndingKey()) {
-                    maxId = TKey::FromBinaryKey(item.GetEndingKey(), Self->Config).GetBlobId();
+                switch (item.GetKeyDesignatorCase()) {
+                    case NKikimrBlobDepot::TEvResolve::TItem::kKeyRange: {
+                        const auto& range = item.GetKeyRange();
+                        if (range.HasBeginningKey()) {
+                            minId = TKey::FromBinaryKey(range.GetBeginningKey(), Self->Config).GetBlobId();
+                        }
+                        if (range.HasEndingKey()) {
+                            maxId = TKey::FromBinaryKey(range.GetEndingKey(), Self->Config).GetBlobId();
+                        }
+                        break;
+                    }
+
+                    case NKikimrBlobDepot::TEvResolve::TItem::kExactKey:
+                        minId = maxId = TKey::FromBinaryKey(item.GetExactKey(), Self->Config).GetBlobId();
+                        break;
+
+                    case NKikimrBlobDepot::TEvResolve::TItem::KEYDESIGNATOR_NOT_SET:
+                        Y_VERIFY_DEBUG(false);
+                        break;
                 }
 
                 Y_VERIFY_DEBUG(minId.TabletID() == tabletId);
diff --git a/ydb/core/protos/blob_depot.proto b/ydb/core/protos/blob_depot.proto
index ce621373296..b24e02726db 100644
--- a/ydb/core/protos/blob_depot.proto
+++ b/ydb/core/protos/blob_depot.proto
@@ -178,18 +178,24 @@ message TEvCommitBlobSeqResult {
 // each of the found entries.
 
 message TEvResolve {
-    message TItem {
+    message TKeyRange {
         optional bytes BeginningKey = 1; // start from the very first key (if not set)
         optional bool IncludeBeginning = 2 [default = true];
         optional bytes EndingKey = 3; // end with the key beyond the last one (if not set)
         optional bool IncludeEnding = 4 [default = false];
         optional uint32 MaxKeys = 5 [default = 0]; // zero or unset value means infinite amount
-        optional bool ReturnMeta = 6 [default = false];
-        optional bool ReturnOwners = 7 [default = false];
-        optional bool Reverse = 8 [default = false]; // reverse output
-        optional uint64 Cookie = 9; // request cookie to match response item
-        optional fixed64 TabletId = 10; // used in virtual group mode to resolve keys of specific tablet
-        optional bool MustRestoreFirst = 11;
+        optional bool Reverse = 6 [default = false]; // reverse output
+    }
+    message TItem {
+        oneof KeyDesignator {
+            TKeyRange KeyRange = 1;
+            bytes ExactKey = 2;
+        }
+        optional bool ReturnMeta = 3 [default = false];
+        optional bool ReturnOwners = 4 [default = false];
+        optional uint64 Cookie = 5; // request cookie to match response item
+        optional fixed64 TabletId = 6; // used in virtual group mode to resolve keys of specific tablet
+        optional bool MustRestoreFirst = 7;
     }
 
     repeated TItem Items = 1;