feat(keyspace): add keyspace notification support for replicas

src/cluster/replication.cc

@@ -30,6 +30,7 @@
 #include <csignal>
 #include <future>
 #include <memory>
+#include <optional>
 #include <string>
 #include <string_view>
 #include <thread>
@@ -40,13 +41,15 @@
 #include "fmt/ostream.h"
 #include "io_util.h"
 #include "logging.h"
+#include "parse_util.h"
 #include "rocksdb/write_batch.h"
 #include "rocksdb_crc32c.h"
 #include "scope_exit.h"
 #include "server/redis_reply.h"
 #include "server/server.h"
 #include "status.h"
 #include "storage/batch_debugger.h"
+#include "storage/redis_db.h"
 #include "thread_util.h"
 #include "time_util.h"
 #include "unique_fd.h"
@@ -57,6 +60,44 @@
 #include <openssl/ssl.h>
 #endif
 
+namespace {
+
+struct ParsedKeyspaceEventContext {
+  KeyspaceEventType type_flag;
+  std::string_view event;
+};
+
+std::optional<ParsedKeyspaceEventContext> ParseKeyspaceEventContext(const rocksdb::Slice &blob) {
+  if (blob.empty() || ServerLogData::IsServerLogData(blob.data())) return std::nullopt;
+
+  redis::WriteBatchLogData log_data;
+  if (!log_data.Decode(blob).IsOK()) return std::nullopt;
+
+  const auto *args = log_data.GetArguments();
+  if (args->empty()) return std::nullopt;
+
+  auto command = ParseInt<int>(args->front(), 10);
+  if (!command) return std::nullopt;
+
+  switch (static_cast<RedisCommand>(*command)) {
+    case kRedisCmdSet:
+      if (log_data.GetRedisType() == kRedisString) {
+        return ParsedKeyspaceEventContext{kNotifyString, "set"};
+      }
+      break;
+    case kRedisCmdDel:
+      if (log_data.GetRedisType() == kRedisNone) {
+        return ParsedKeyspaceEventContext{kNotifyGeneric, "del"};
+      }
+      break;
+    default:
+      break;
+  }
+  return std::nullopt;
+}
+
+}  // namespace
+
 FeedSlaveThread::FeedSlaveThread(Server *srv, redis::Connection *conn, rocksdb::SequenceNumber next_repl_seq)
     : srv_(srv),
       conn_(conn),
@@ -1155,7 +1196,10 @@ void ReplicationThread::TimerCB(int, int16_t) {
 }
 
 Status ReplicationThread::parseWriteBatch(const rocksdb::WriteBatch &write_batch) {
-  WriteBatchHandler write_batch_handler;
+  const auto *config = srv_->GetConfig();
+  const bool keyspace_notifications_enabled = config->notify_keyspace_event_channels != kNotifyNoChannel &&
+                                              (config->notify_keyspace_event_types & kNotifyAll) != 0;
+  WriteBatchHandler write_batch_handler(keyspace_notifications_enabled);
 
   auto db_status = write_batch.Iterate(&write_batch_handler);
   if (!db_status.ok()) return {Status::NotOK, "failed to iterate over write batch: " + db_status.ToString()};
@@ -1193,6 +1237,20 @@ Status ReplicationThread::parseWriteBatch(const rocksdb::WriteBatch &write_batch
     case kBatchTypeNone:
       break;
   }
+
+  if (keyspace_notifications_enabled && write_batch_handler.HasKeyspaceEvents()) {
+    KeyspaceEventBatchHandler keyspace_event_handler(storage_->IsSlotIdEncoded());
+    db_status = write_batch.Iterate(&keyspace_event_handler);
+    if (!db_status.ok()) {
+      WARN("[notify] failed to inspect replicated batch for keyspace notifications: {}", db_status.ToString());
+    } else {
+      for (const auto &event : keyspace_event_handler.Events()) {
+        if ((config->notify_keyspace_event_types & event.type_flag) == 0) continue;
+        srv_->NotifyKeyspaceEvent(
+            KeyspaceEvent(event.type_flag, event.event, config->notify_keyspace_event_channels, event.ns, event.key));
+      }
+    }
+  }
   return Status::OK();
 }
 
@@ -1229,3 +1287,57 @@ rocksdb::Status WriteBatchHandler::PutCF(uint32_t column_family_id, const rocksd
   }
   return rocksdb::Status::OK();
 }
+
+void WriteBatchHandler::LogData(const rocksdb::Slice &blob) {
+  if (detect_keyspace_events_ && ParseKeyspaceEventContext(blob)) has_keyspace_events_ = true;
+}
+
+void KeyspaceEventBatchHandler::LogData(const rocksdb::Slice &blob) {
+  current_type_flag_ = kNotifyNoType;
+  current_event_ = {};
+
+  auto event_context = ParseKeyspaceEventContext(blob);
+  if (!event_context) return;
+
+  current_type_flag_ = event_context->type_flag;
+  current_event_ = event_context->event;
+}
+
+rocksdb::Status KeyspaceEventBatchHandler::PutCF(uint32_t column_family_id, const rocksdb::Slice &key,
+                                                 const rocksdb::Slice &value) {
+  if (current_event_ == "set") {
+    return handleSet(column_family_id, key, value);
+  } else {
+    return rocksdb::Status::OK();
+  }
+}
+
+rocksdb::Status KeyspaceEventBatchHandler::DeleteCF(uint32_t column_family_id, const rocksdb::Slice &key) {
+  if (current_event_ == "del") {
+    return handleDel(column_family_id, key);
+  } else {
+    return rocksdb::Status::OK();
+  }
+}
+
+rocksdb::Status KeyspaceEventBatchHandler::handleSet(uint32_t column_family_id, const rocksdb::Slice &key,
+                                                     const rocksdb::Slice &value) {
+  if (column_family_id != static_cast<uint32_t>(ColumnFamilyID::Metadata)) return rocksdb::Status::OK();
+
+  Metadata metadata(kRedisNone, false);
+  if (auto s = metadata.Decode(value); !s.ok() || metadata.Type() != kRedisString) {
+    return rocksdb::Status::OK();
+  }
+
+  auto [ns, user_key] = ExtractNamespaceKey<std::string>(key, is_slot_id_encoded_);
+  keyspace_events_.emplace_back(current_type_flag_, current_event_, kNotifyNoChannel, ns, user_key);
+  return rocksdb::Status::OK();
+}
+
+rocksdb::Status KeyspaceEventBatchHandler::handleDel(uint32_t column_family_id, const rocksdb::Slice &key) {
+  if (column_family_id != static_cast<uint32_t>(ColumnFamilyID::Metadata)) return rocksdb::Status::OK();
+
+  auto [ns, user_key] = ExtractNamespaceKey<std::string>(key, is_slot_id_encoded_);
+  keyspace_events_.emplace_back(current_type_flag_, current_event_, kNotifyNoChannel, ns, user_key);
+  return rocksdb::Status::OK();
+}

src/cluster/replication.h

@@ -27,6 +27,7 @@
 #include <deque>
 #include <memory>
 #include <string>
+#include <string_view>
 #include <thread>
 #include <tuple>
 #include <utility>
@@ -230,6 +231,8 @@ class ReplicationThread : private EventCallbackBase<ReplicationThread> {
  */
 class WriteBatchHandler : public rocksdb::WriteBatch::Handler {
  public:
+  explicit WriteBatchHandler(bool detect_keyspace_events = false) : detect_keyspace_events_(detect_keyspace_events) {}
+
   rocksdb::Status PutCF(uint32_t column_family_id, const rocksdb::Slice &key, const rocksdb::Slice &value) override;
   rocksdb::Status DeleteCF([[maybe_unused]] uint32_t column_family_id,
                            [[maybe_unused]] const rocksdb::Slice &key) override {
@@ -240,11 +243,39 @@ class WriteBatchHandler : public rocksdb::WriteBatch::Handler {
                                 [[maybe_unused]] const rocksdb::Slice &end_key) override {
     return rocksdb::Status::OK();
   }
+  void LogData(const rocksdb::Slice &blob) override;
   WriteBatchType Type() { return type_; }
   std::string Key() const { return kv_.first; }
   std::string Value() const { return kv_.second; }
+  bool HasKeyspaceEvents() const { return has_keyspace_events_; }
 
  private:
   std::pair<std::string, std::string> kv_;
   WriteBatchType type_ = kBatchTypeNone;
+  bool detect_keyspace_events_ = false;
+  bool has_keyspace_events_ = false;
+};
+
+class KeyspaceEventBatchHandler : public rocksdb::WriteBatch::Handler {
+ public:
+  explicit KeyspaceEventBatchHandler(bool is_slot_id_encoded) : is_slot_id_encoded_(is_slot_id_encoded) {}
+
+  rocksdb::Status PutCF(uint32_t column_family_id, const rocksdb::Slice &key, const rocksdb::Slice &value) override;
+  rocksdb::Status DeleteCF(uint32_t column_family_id, const rocksdb::Slice &key) override;
+  rocksdb::Status DeleteRangeCF([[maybe_unused]] uint32_t column_family_id,
+                                [[maybe_unused]] const rocksdb::Slice &begin_key,
+                                [[maybe_unused]] const rocksdb::Slice &end_key) override {
+    return rocksdb::Status::OK();
+  }
+  void LogData(const rocksdb::Slice &blob) override;
+  const std::vector<KeyspaceEvent> &Events() const { return keyspace_events_; }
+
+ private:
+  rocksdb::Status handleSet(uint32_t column_family_id, const rocksdb::Slice &key, const rocksdb::Slice &value);
+  rocksdb::Status handleDel(uint32_t column_family_id, const rocksdb::Slice &key);
+
+  bool is_slot_id_encoded_ = false;
+  KeyspaceEventType current_type_flag_ = kNotifyNoType;
+  std::string_view current_event_;
+  std::vector<KeyspaceEvent> keyspace_events_;
 };

src/storage/redis_db.cc

@@ -171,7 +171,7 @@ rocksdb::Status Database::MDel(engine::Context &ctx, const std::vector<Slice> &k
   }
 
   auto batch = storage_->GetWriteBatchBase();
-  WriteBatchLogData log_data(kRedisNone);
+  WriteBatchLogData log_data(kRedisNone, {std::to_string(kRedisCmdDel)});
   auto s = batch->PutLogData(log_data.Encode());
   if (!s.ok()) {
     return s;

src/storage/redis_metadata.h

@@ -99,6 +99,8 @@ enum RedisCommand {
   kRedisCmdBitOp,
   kRedisCmdBitfield,
   kRedisCmdLMove,
+  kRedisCmdSet,
+  kRedisCmdDel,
 };
 
 constexpr const char *kErrMsgWrongType = "WRONGTYPE Operation against a key holding the wrong kind of value";

src/types/redis_string.cc

@@ -107,9 +107,11 @@ std::vector<rocksdb::Status> String::getValues(engine::Context &ctx, const std::
   return statuses;
 }
 
-rocksdb::Status String::updateRawValue(engine::Context &ctx, const std::string &ns_key, const std::string &raw_value) {
+rocksdb::Status String::updateRawValue(engine::Context &ctx, const std::string &ns_key, const std::string &raw_value,
+                                       std::optional<RedisCommand> command) {
   auto batch = storage_->GetWriteBatchBase();
   WriteBatchLogData log_data(kRedisString);
+  if (command) log_data.GetArguments()->emplace_back(std::to_string(*command));
   auto s = batch->PutLogData(log_data.Encode());
   if (!s.ok()) return s;
   s = batch->Put(metadata_cf_handle_, ns_key, raw_value);
@@ -345,7 +347,7 @@ rocksdb::Status String::Set(engine::Context &ctx, const std::string &user_key, c
   metadata.expire = expire;
   metadata.Encode(&new_raw_value);
   new_raw_value.append(value);
-  auto s = updateRawValue(ctx, ns_key, new_raw_value);
+  auto s = updateRawValue(ctx, ns_key, new_raw_value, kRedisCmdSet);
   if (!s.ok()) return s;
 
   ctx.AddKeyspaceEventIfEnabled(kNotifyString, "set", namespace_, user_key);

src/types/redis_string.h

@@ -138,7 +138,8 @@ class String : public Database {
   rocksdb::Status getRawValue(engine::Context &ctx, const std::string &ns_key, std::string *raw_value);
   std::vector<rocksdb::Status> getRawValues(engine::Context &ctx, const std::vector<Slice> &keys,
                                             std::vector<std::string> *raw_values);
-  rocksdb::Status updateRawValue(engine::Context &ctx, const std::string &ns_key, const std::string &raw_value);
+  rocksdb::Status updateRawValue(engine::Context &ctx, const std::string &ns_key, const std::string &raw_value,
+                                 std::optional<RedisCommand> command = std::nullopt);
 };
 
 }  // namespace redis

tests/cppunit/replication_test.cc

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include "cluster/replication.h"
+
+#include <gtest/gtest.h>
+
+#include "storage/redis_db.h"
+#include "storage/redis_metadata.h"
+#include "test_base.h"
+
+class ReplicationWriteBatchTest : public TestBase {};
+
+TEST_F(ReplicationWriteBatchTest, ExtractKeyspaceEvents) {
+  rocksdb::WriteBatch batch;
+  auto *metadata_cf = storage_->GetCFHandle(ColumnFamilyID::Metadata);
+
+  redis::WriteBatchLogData set_log_data(kRedisString, {std::to_string(kRedisCmdSet)});
+  ASSERT_TRUE(batch.PutLogData(set_log_data.Encode()).ok());
+  std::string string_value;
+  Metadata(kRedisString, false).Encode(&string_value);
+  string_value.append("value");
+  ASSERT_TRUE(batch.Put(metadata_cf, ComposeNamespaceKey("tenant", "set-key", false), string_value).ok());
+
+  redis::WriteBatchLogData legacy_log_data(kRedisString);
+  ASSERT_TRUE(batch.PutLogData(legacy_log_data.Encode()).ok());
+  ASSERT_TRUE(batch.Put(metadata_cf, ComposeNamespaceKey("tenant", "legacy-key", false), string_value).ok());
+
+  redis::WriteBatchLogData del_log_data(kRedisNone, {std::to_string(kRedisCmdDel)});
+  ASSERT_TRUE(batch.PutLogData(del_log_data.Encode()).ok());
+  ASSERT_TRUE(batch.Delete(metadata_cf, ComposeNamespaceKey("tenant", "del-key", false)).ok());
+
+  WriteBatchHandler detector(true);
+  ASSERT_TRUE(batch.Iterate(&detector).ok());
+  ASSERT_TRUE(detector.HasKeyspaceEvents());
+
+  KeyspaceEventBatchHandler handler(false);
+  ASSERT_TRUE(batch.Iterate(&handler).ok());
+  const auto &events = handler.Events();
+  ASSERT_EQ(events.size(), 2);
+  EXPECT_EQ(events[0].type_flag, kNotifyString);
+  EXPECT_EQ(events[0].event, "set");
+  EXPECT_EQ(events[0].ns, "tenant");
+  EXPECT_EQ(events[0].key, "set-key");
+  EXPECT_EQ(events[1].type_flag, kNotifyGeneric);
+  EXPECT_EQ(events[1].event, "del");
+  EXPECT_EQ(events[1].ns, "tenant");
+  EXPECT_EQ(events[1].key, "del-key");
+}
+
+TEST_F(ReplicationWriteBatchTest, IgnoreLegacyLogData) {
+  rocksdb::WriteBatch batch;
+  redis::WriteBatchLogData legacy_log_data(kRedisString);
+  ASSERT_TRUE(batch.PutLogData(legacy_log_data.Encode()).ok());
+
+  std::string string_value;
+  Metadata(kRedisString, false).Encode(&string_value);
+  string_value.append("value");
+  ASSERT_TRUE(batch
+                  .Put(storage_->GetCFHandle(ColumnFamilyID::Metadata),
+                       ComposeNamespaceKey(kDefaultNamespace, "key", false), string_value)
+                  .ok());
+
+  WriteBatchHandler detector(true);
+  ASSERT_TRUE(batch.Iterate(&detector).ok());
+  EXPECT_FALSE(detector.HasKeyspaceEvents());
+
+  KeyspaceEventBatchHandler handler(false);
+  ASSERT_TRUE(batch.Iterate(&handler).ok());
+  EXPECT_TRUE(handler.Events().empty());
+}

tests/gocase/unit/keyspacenotify/keyspacenotify_test.go

@@ -233,6 +233,36 @@ func TestKeyspaceNotifyDisabled(t *testing.T) {
 	require.Error(t, rdb.ConfigSet(ctx, "notify-keyspace-events", "KEl").Err())
 }
 
+func TestKeyspaceNotifyOnReplica(t *testing.T) {
+	master := util.StartServer(t, map[string]string{})
+	defer master.Close()
+	masterClient := master.NewClient()
+	defer func() { require.NoError(t, masterClient.Close()) }()
+
+	replica := util.StartServer(t, map[string]string{"notify-keyspace-events": "KEA"})
+	defer replica.Close()
+	replicaClient := replica.NewClient()
+	defer func() { require.NoError(t, replicaClient.Close()) }()
+
+	ctx := context.Background()
+	util.SlaveOf(t, replicaClient, master)
+	util.WaitForSync(t, replicaClient)
+
+	subscriber := replica.NewClient()
+	defer func() { require.NoError(t, subscriber.Close()) }()
+	pubsub := subscriber.PSubscribe(ctx, "__keyspace@0__:*", "__keyevent@0__:*")
+	defer func() { require.NoError(t, pubsub.Close()) }()
+	drainSubscribeConfirms(t, ctx, pubsub, 2)
+
+	require.NoError(t, masterClient.Set(ctx, "replicated-key", "value", 0).Err())
+	expectMessage(t, ctx, pubsub, "__keyspace@0__:replicated-key", "set")
+	expectMessage(t, ctx, pubsub, "__keyevent@0__:set", "replicated-key")
+
+	require.EqualValues(t, 1, masterClient.Del(ctx, "replicated-key").Val())
+	expectMessage(t, ctx, pubsub, "__keyspace@0__:replicated-key", "del")
+	expectMessage(t, ctx, pubsub, "__keyevent@0__:del", "replicated-key")
+}
+
 func TestKeyspaceNotifyRedisDatabases(t *testing.T) {
 	srv := util.StartServer(t, map[string]string{"notify-keyspace-events": "KEA", "redis-databases": "16"})
 	defer srv.Close()