Add support for Put with secondary indices (facebook#13289)

Summary:

The patch adds support for `Put` / `PutUntracked` to the secondary indexing logic. Similarly to `PutEntity` (see facebook#13180), calling these APIs automatically add or remove secondary index entries as needed in an atomic and transparent fashion.

Differential Revision: D68035089

utilities/secondary_index/secondary_index_mixin.h

@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include <array>
 #include <cassert>
 #include <memory>
 #include <optional>
@@ -35,15 +36,23 @@ class SecondaryIndexMixin : public Txn {
   }
 
   using Txn::Put;
-  Status Put(ColumnFamilyHandle* /* column_family */, const Slice& /* key */,
-             const Slice& /* value */,
-             const bool /* assume_tracked */ = false) override {
-    return Status::NotSupported("Put with secondary indices not yet supported");
+  Status Put(ColumnFamilyHandle* column_family, const Slice& key,
+             const Slice& value, const bool assume_tracked = false) override {
+    return PerformWithSavePoint([&]() {
+      const bool do_validate = !assume_tracked;
+      return PutWithSecondaryIndices(column_family, key, value, do_validate);
+    });
   }
-  Status Put(ColumnFamilyHandle* /* column_family */,
-             const SliceParts& /* key */, const SliceParts& /* value */,
-             const bool /* assume_tracked */ = false) override {
-    return Status::NotSupported("Put with secondary indices not yet supported");
+  Status Put(ColumnFamilyHandle* column_family, const SliceParts& key,
+             const SliceParts& value,
+             const bool assume_tracked = false) override {
+    std::string key_str;
+    const Slice key_slice(key, &key_str);
+
+    std::string value_str;
+    const Slice value_slice(value, &value_str);
+
+    return Put(column_family, key_slice, value_slice, assume_tracked);
   }
 
   Status PutEntity(ColumnFamilyHandle* column_family, const Slice& key,
@@ -92,17 +101,22 @@ class SecondaryIndexMixin : public Txn {
   }
 
   using Txn::PutUntracked;
-  Status PutUntracked(ColumnFamilyHandle* /* column_family */,
-                      const Slice& /* key */,
-                      const Slice& /* value */) override {
-    return Status::NotSupported(
-        "PutUntracked with secondary indices not yet supported");
+  Status PutUntracked(ColumnFamilyHandle* column_family, const Slice& key,
+                      const Slice& value) override {
+    return PerformWithSavePoint([&]() {
+      constexpr bool do_validate = false;
+      return PutWithSecondaryIndices(column_family, key, value, do_validate);
+    });
   }
-  Status PutUntracked(ColumnFamilyHandle* /* column_family */,
-                      const SliceParts& /* key */,
-                      const SliceParts& /* value */) override {
-    return Status::NotSupported(
-        "PutUntracked with secondary indices not yet supported");
+  Status PutUntracked(ColumnFamilyHandle* column_family, const SliceParts& key,
+                      const SliceParts& value) override {
+    std::string key_str;
+    const Slice key_slice(key, &key_str);
+
+    std::string value_str;
+    const Slice value_slice(value, &value_str);
+
+    return PutUntracked(column_family, key_slice, value_slice);
   }
 
   Status PutEntityUntracked(ColumnFamilyHandle* column_family, const Slice& key,
@@ -194,8 +208,7 @@ class SecondaryIndexMixin : public Txn {
                                    secondary_index->GetPrimaryColumnName());
   }
 
-  template <typename Operation>
-  Status PerformWithSavePoint(Operation&& operation) {
+  Status PerformWithSavePoint(auto&& operation) {
     Txn::SetSavePoint();
 
     const Status s = operation();
@@ -247,17 +260,6 @@ class SecondaryIndexMixin : public Txn {
                              secondary_key);
   }
 
-  Status AddPrimaryEntry(ColumnFamilyHandle* column_family,
-                         const Slice& primary_key,
-                         const WideColumns& primary_columns) {
-    assert(column_family);
-
-    constexpr bool assume_tracked = true;
-
-    return Txn::PutEntity(column_family, primary_key, primary_columns,
-                          assume_tracked);
-  }
-
   Status AddSecondaryEntry(const SecondaryIndex* secondary_index,
                            const Slice& primary_key,
                            const Slice& primary_column_value,
@@ -336,7 +338,7 @@ class SecondaryIndexMixin : public Txn {
 
   Status UpdatePrimaryColumnValues(ColumnFamilyHandle* column_family,
                                    const Slice& primary_key,
-                                   WideColumns& primary_columns,
+                                   auto& primary_columns,
                                    autovector<IndexData>& applicable_indices) {
     assert(applicable_indices.empty());
 
@@ -382,10 +384,11 @@ class SecondaryIndexMixin : public Txn {
     return Status::OK();
   }
 
-  Status PutEntityWithSecondaryIndices(ColumnFamilyHandle* column_family,
-                                       const Slice& key,
-                                       const WideColumns& columns,
-                                       bool do_validate) {
+  Status PutlikeWithSecondaryIndices(ColumnFamilyHandle* column_family,
+                                     const auto& key, const auto& valuelike,
+                                     bool do_validate,
+                                     auto&& get_primary_columns,
+                                     auto&& add_primary_entry) {
     // TODO: we could avoid removing and recreating secondary entries for
     // which neither the secondary key prefix nor the value has changed
 
@@ -403,11 +406,9 @@ class SecondaryIndexMixin : public Txn {
       }
     }
 
+    auto primary_columns = get_primary_columns(valuelike);
     autovector<IndexData> applicable_indices;
 
-    WideColumns primary_columns(columns);
-    WideColumnsHelper::SortColumns(primary_columns);
-
     {
       const Status s = UpdatePrimaryColumnValues(
           column_family, primary_key, primary_columns, applicable_indices);
@@ -418,7 +419,7 @@ class SecondaryIndexMixin : public Txn {
 
     {
       const Status s =
-          AddPrimaryEntry(column_family, primary_key, primary_columns);
+          add_primary_entry(column_family, primary_key, primary_columns);
       if (!s.ok()) {
         return s;
       }
@@ -434,6 +435,50 @@ class SecondaryIndexMixin : public Txn {
     return Status::OK();
   }
 
+  Status PutWithSecondaryIndices(ColumnFamilyHandle* column_family,
+                                 const Slice& key, const Slice& value,
+                                 bool do_validate) {
+    return PutlikeWithSecondaryIndices(
+        column_family, key, value, do_validate,
+        [](const Slice& v) {
+          return std::array<WideColumn, 1>{{{kDefaultWideColumnName, v}}};
+        },
+        [this](ColumnFamilyHandle* cfh, const Slice& primary_key,
+               const std::array<WideColumn, 1>& primary_columns) {
+          assert(cfh);
+          assert(!primary_columns.empty());
+          assert(primary_columns.front().name() == kDefaultWideColumnName);
+
+          constexpr bool assume_tracked = true;
+
+          return Txn::Put(cfh, primary_key, primary_columns.front().value(),
+                          assume_tracked);
+        });
+  }
+
+  Status PutEntityWithSecondaryIndices(ColumnFamilyHandle* column_family,
+                                       const Slice& key,
+                                       const WideColumns& columns,
+                                       bool do_validate) {
+    return PutlikeWithSecondaryIndices(
+        column_family, key, columns, do_validate,
+        [](const WideColumns& c) {
+          WideColumns primary_columns(c);
+          WideColumnsHelper::SortColumns(primary_columns);
+
+          return primary_columns;
+        },
+        [this](ColumnFamilyHandle* cfh, const Slice& primary_key,
+               const WideColumns& primary_columns) {
+          assert(cfh);
+
+          constexpr bool assume_tracked = true;
+
+          return Txn::PutEntity(cfh, primary_key, primary_columns,
+                                assume_tracked);
+        });
+  }
+
   const std::vector<std::shared_ptr<SecondaryIndex>>* secondary_indices_;
 };