libroach: add basic encryption stats.

Makefile

@@ -1017,7 +1017,7 @@ CPP_PROTOS := $(filter %/roachpb/metadata.proto %/roachpb/data.proto %/roachpb/i
 CPP_HEADERS := $(subst $(PKG_ROOT),$(CPP_PROTO_ROOT),$(CPP_PROTOS:%.proto=%.pb.h))
 CPP_SOURCES := $(subst $(PKG_ROOT),$(CPP_PROTO_ROOT),$(CPP_PROTOS:%.proto=%.pb.cc))
 
-CPP_PROTOS_CCL := $(filter %/ccl/baseccl/encryption_options.proto %/ccl/storageccl/engineccl/enginepbccl/key_registry.proto,$(GO_PROTOS))
+CPP_PROTOS_CCL := $(filter %/ccl/baseccl/encryption_options.proto %/ccl/storageccl/engineccl/enginepbccl/key_registry.proto %/ccl/storageccl/engineccl/enginepbccl/stats.proto,$(GO_PROTOS))
 CPP_HEADERS_CCL := $(subst $(PKG_ROOT),$(CPP_PROTO_CCL_ROOT),$(CPP_PROTOS_CCL:%.proto=%.pb.h))
 CPP_SOURCES_CCL := $(subst $(PKG_ROOT),$(CPP_PROTO_CCL_ROOT),$(CPP_PROTOS_CCL:%.proto=%.pb.cc))
 

c-deps/libroach/CMakeLists.txt

@@ -65,6 +65,7 @@ add_library(roachccl
   ccl/key_manager.cc
   protosccl/ccl/baseccl/encryption_options.pb.cc
   protosccl/ccl/storageccl/engineccl/enginepbccl/key_registry.pb.cc
+  protosccl/ccl/storageccl/engineccl/enginepbccl/stats.pb.cc
 )
 target_include_directories(roachccl
   PRIVATE .. # CryptoPP headers are directly in the directory. Include .. to be able to include <cryptopp/....h>

c-deps/libroach/batch.cc

@@ -513,6 +513,8 @@ DBStatus DBBatch::GetStats(DBStatsResult* stats) { return FmtStatus("unsupported
 
 DBString DBBatch::GetCompactionStats() { return ToDBString("unsupported"); }
 
+DBStatus DBBatch::GetEnvStats(DBEnvStatsResult* stats) { return FmtStatus("unsupported"); }
+
 DBStatus DBBatch::EnvWriteFile(DBSlice path, DBSlice contents) { return FmtStatus("unsupported"); }
 
 DBStatus DBBatch::EnvOpenFile(DBSlice path, rocksdb::WritableFile** file) {
@@ -590,6 +592,8 @@ DBStatus DBWriteOnlyBatch::GetStats(DBStatsResult* stats) { return FmtStatus("un
 
 DBString DBWriteOnlyBatch::GetCompactionStats() { return ToDBString("unsupported"); }
 
+DBStatus DBWriteOnlyBatch::GetEnvStats(DBEnvStatsResult* stats) { return FmtStatus("unsupported"); }
+
 DBStatus DBWriteOnlyBatch::EnvWriteFile(DBSlice path, DBSlice contents) {
   return FmtStatus("unsupported");
 }

c-deps/libroach/batch.h

@@ -40,6 +40,7 @@ struct DBBatch : public DBEngine {
   virtual DBIterator* NewIter(rocksdb::ReadOptions*);
   virtual DBStatus GetStats(DBStatsResult* stats);
   virtual DBString GetCompactionStats();
+  virtual DBStatus GetEnvStats(DBEnvStatsResult* stats);
   virtual DBStatus EnvWriteFile(DBSlice path, DBSlice contents);
   virtual DBStatus EnvOpenFile(DBSlice path, rocksdb::WritableFile** file);
   virtual DBStatus EnvAppendFile(rocksdb::WritableFile* file, DBSlice contents);
@@ -65,6 +66,7 @@ struct DBWriteOnlyBatch : public DBEngine {
   virtual DBIterator* NewIter(rocksdb::ReadOptions*);
   virtual DBStatus GetStats(DBStatsResult* stats);
   virtual DBString GetCompactionStats();
+  virtual DBString GetEnvStats(DBEnvStatsResult* stats);
   virtual DBStatus EnvWriteFile(DBSlice path, DBSlice contents);
   virtual DBStatus EnvOpenFile(DBSlice path, rocksdb::WritableFile** file);
   virtual DBStatus EnvAppendFile(rocksdb::WritableFile* file, DBSlice contents);

c-deps/libroach/ccl/db.cc

@@ -21,13 +21,52 @@
 #include "../rocksdbutils/env_encryption.h"
 #include "../status.h"
 #include "ccl/baseccl/encryption_options.pb.h"
+#include "ccl/storageccl/engineccl/enginepbccl/stats.pb.h"
 #include "ctr_stream.h"
 #include "key_manager.h"
 
 using namespace cockroach;
 
 namespace cockroach {
 
+class CCLEnvStatsHandler : public EnvStatsHandler {
+ public:
+  explicit CCLEnvStatsHandler(KeyManager* store_key_manager, KeyManager* data_key_manager)
+      : store_key_manager_(store_key_manager), data_key_manager_(data_key_manager) {}
+  virtual ~CCLEnvStatsHandler() {}
+
+  virtual rocksdb::Status GetEncryptionStats(std::string* serialized_stats) override {
+    enginepbccl::EncryptionStatus enc_status;
+
+    bool has_stats = false;
+    if (this->store_key_manager_ != nullptr) {
+      has_stats = true;
+      // Transfer ownership of new key info to status proto, this frees any previous value.
+      enc_status.set_allocated_active_store_key(store_key_manager_->CurrentKeyInfo().release());
+    }
+
+    if (this->data_key_manager_ != nullptr) {
+      has_stats = true;
+      // Transfer ownership of new key info to status proto, this frees any previous value.
+      enc_status.set_allocated_active_data_key(data_key_manager_->CurrentKeyInfo().release());
+    }
+
+    if (!has_stats) {
+      return rocksdb::Status::OK();
+    }
+
+    if (!enc_status.SerializeToString(serialized_stats)) {
+      return rocksdb::Status::InvalidArgument("failed to serialize encryption status");
+    }
+    return rocksdb::Status::OK();
+  }
+
+ private:
+  // KeyManagers are needed to get key information but are not owned by the StatsHandler.
+  KeyManager* store_key_manager_;
+  KeyManager* data_key_manager_;
+};
+
 // DBOpenHook parses the extra_options field of DBOptions and initializes encryption objects if
 // needed.
 rocksdb::Status DBOpenHook(const std::string& db_dir, const DBOptions db_opts,
@@ -98,12 +137,16 @@ rocksdb::Status DBOpenHook(const std::string& db_dir, const DBOptions db_opts,
   std::unique_ptr<enginepbccl::KeyInfo> store_key = store_key_manager->CurrentKeyInfo();
   assert(store_key != nullptr);
 
-  // Generate a new data key if needed by giving the active store key info to the data key manager.
+  // Generate a new data key if needed by giving the active store key info to the data key
+  // manager.
   status = data_key_manager->SetActiveStoreKey(std::move(store_key));
   if (!status.ok()) {
     return status;
   }
 
+  // Everything's ok: initialize a stats handler.
+  env_mgr->SetStatsHandler(new CCLEnvStatsHandler(store_key_manager, data_key_manager));
+
   return rocksdb::Status::OK();
 }
 

c-deps/libroach/ccl/db_test.cc

@@ -8,6 +8,8 @@
 
 #include "../db.h"
 #include "../testutils.h"
+#include "ccl/baseccl/encryption_options.pb.h"
+#include "ccl/storageccl/engineccl/enginepbccl/stats.pb.h"
 
 using namespace cockroach;
 
@@ -30,3 +32,49 @@ TEST(LibroachCCL, DBOpenHook) {
   db_opts.extra_options = ToDBSlice("blah");
   EXPECT_ERR(DBOpenHook("", db_opts, nullptr), "failed to parse extra options");
 }
+
+TEST(Libroach, DBOpen) {
+  {
+    // Empty options: no encryption.
+    DBOptions db_opts = defaultDBOptions();
+    DBEngine* db;
+    db_opts.use_file_registry = true;
+
+    EXPECT_STREQ(DBOpen(&db, DBSlice(), db_opts).data, NULL);
+    DBEnvStatsResult stats;
+    EXPECT_STREQ(DBGetEnvStats(db, &stats).data, NULL);
+    EXPECT_STREQ(stats.encryption_status.data, NULL);
+
+    DBClose(db);
+  }
+  {
+    // Encryption enabled.
+    DBOptions db_opts = defaultDBOptions();
+    DBEngine* db;
+    db_opts.use_file_registry = true;
+
+    // Enable encryption, but plaintext only, that's enough to get stats going.
+    cockroach::ccl::baseccl::EncryptionOptions enc_opts;
+    enc_opts.set_key_source(cockroach::ccl::baseccl::KeyFiles);
+    enc_opts.mutable_key_files()->set_current_key("plain");
+    enc_opts.mutable_key_files()->set_old_key("plain");
+
+    std::string tmpstr;
+    ASSERT_TRUE(enc_opts.SerializeToString(&tmpstr));
+    db_opts.extra_options = ToDBSlice(tmpstr);
+
+    EXPECT_STREQ(DBOpen(&db, DBSlice(), db_opts).data, NULL);
+    DBEnvStatsResult stats;
+    EXPECT_STREQ(DBGetEnvStats(db, &stats).data, NULL);
+    EXPECT_STRNE(stats.encryption_status.data, NULL);
+
+    // Now parse the status protobuf.
+    cockroach::ccl::storageccl::engineccl::enginepbccl::EncryptionStatus enc_status;
+    ASSERT_TRUE(
+        enc_status.ParseFromArray(stats.encryption_status.data, stats.encryption_status.len));
+    EXPECT_STREQ(enc_status.active_store_key().key_id().c_str(), "plain");
+    EXPECT_STREQ(enc_status.active_data_key().key_id().c_str(), "plain");
+
+    DBClose(db);
+  }
+}

c-deps/libroach/db.cc

@@ -628,6 +628,8 @@ DBStatus DBGetStats(DBEngine* db, DBStatsResult* stats) { return db->GetStats(st
 
 DBString DBGetCompactionStats(DBEngine* db) { return db->GetCompactionStats(); }
 
+DBStatus DBGetEnvStats(DBEngine* db, DBEnvStatsResult* stats) { return db->GetEnvStats(stats); }
+
 DBSSTable* DBGetSSTables(DBEngine* db, int* n) { return db->GetSSTables(n); }
 
 DBString DBGetUserProperties(DBEngine* db) { return db->GetUserProperties(); }

c-deps/libroach/db_test.cc

@@ -14,6 +14,7 @@
 
 #include "db.h"
 #include "include/libroach.h"
+#include "status.h"
 #include "testutils.h"
 
 using namespace cockroach;
@@ -30,3 +31,15 @@ TEST(Libroach, DBOpenHook) {
   EXPECT_ERR(DBOpenHook("", db_opts, nullptr),
              "DBOptions has extra_options, but OSS code cannot handle them");
 }
+
+TEST(Libroach, DBOpen) {
+  DBOptions db_opts = defaultDBOptions();
+  DBEngine* db;
+
+  EXPECT_STREQ(DBOpen(&db, DBSlice(), db_opts).data, NULL);
+  DBEnvStatsResult stats;
+  EXPECT_STREQ(DBGetEnvStats(db, &stats).data, NULL);
+  EXPECT_STREQ(stats.encryption_status.data, NULL);
+
+  DBClose(db);
+}

c-deps/libroach/engine.cc

@@ -204,6 +204,24 @@ DBString DBImpl::GetCompactionStats() {
   return ToDBString(tmp);
 }
 
+DBStatus DBImpl::GetEnvStats(DBEnvStatsResult* stats) {
+  // Always initialize the field to the empty string.
+  stats->encryption_status = DBString();
+
+  if (this->env_mgr->env_stats_handler == nullptr) {
+    return kSuccess;
+  }
+
+  std::string encryption_status;
+  auto status = this->env_mgr->env_stats_handler->GetEncryptionStats(&encryption_status);
+  if (!status.ok()) {
+    return ToDBStatus(status);
+  }
+
+  stats->encryption_status = ToDBString(encryption_status);
+  return kSuccess;
+}
+
 // EnvWriteFile writes the given data as a new "file" in the given engine.
 DBStatus DBImpl::EnvWriteFile(DBSlice path, DBSlice contents) {
   rocksdb::Status s;

c-deps/libroach/engine.h

@@ -41,6 +41,7 @@ struct DBEngine {
   virtual DBIterator* NewIter(rocksdb::ReadOptions*) = 0;
   virtual DBStatus GetStats(DBStatsResult* stats) = 0;
   virtual DBString GetCompactionStats() = 0;
+  virtual DBString GetEnvStats(DBEnvStatsResult* stats) = 0;
   virtual DBStatus EnvWriteFile(DBSlice path, DBSlice contents) = 0;
   virtual DBStatus EnvOpenFile(DBSlice path, rocksdb::WritableFile** file) = 0;
   virtual DBStatus EnvAppendFile(rocksdb::WritableFile* file, DBSlice contents) = 0;
@@ -81,6 +82,7 @@ struct DBImpl : public DBEngine {
   virtual DBIterator* NewIter(rocksdb::ReadOptions*);
   virtual DBStatus GetStats(DBStatsResult* stats);
   virtual DBString GetCompactionStats();
+  virtual DBStatus GetEnvStats(DBEnvStatsResult* stats);
   virtual DBStatus EnvWriteFile(DBSlice path, DBSlice contents);
   virtual DBStatus EnvOpenFile(DBSlice path, rocksdb::WritableFile** file);
   virtual DBStatus EnvAppendFile(rocksdb::WritableFile* file, DBSlice contents);

c-deps/libroach/env_manager.h

@@ -20,6 +20,14 @@
 
 namespace cockroach {
 
+// EnvStatsHandler provides an interface to generate Env-specific stats.
+class EnvStatsHandler {
+ public:
+  virtual ~EnvStatsHandler() {}
+
+  virtual rocksdb::Status GetEncryptionStats(std::string* stats) = 0;
+};
+
 // EnvManager manages all created Envs, as well as the file registry.
 // Rocksdb owns Env::Default (global static). All other envs are owned by EnvManager.
 //
@@ -30,10 +38,14 @@ struct EnvManager {
   EnvManager(rocksdb::Env* env) : base_env(env), db_env(env) {}
   ~EnvManager() {}
 
+  // Set the stats handler implementing GetEncryptionStats to fill in env-related stats.
+  // It does not have called, leaving env_stats_handler nil.
+  void SetStatsHandler(EnvStatsHandler* stats_handler) { env_stats_handler.reset(stats_handler); }
   void TakeEnvOwnership(rocksdb::Env* env) { envs.push_back(std::unique_ptr<rocksdb::Env>(env)); }
 
   rocksdb::Env* base_env;
   rocksdb::Env* db_env;
+  std::unique_ptr<EnvStatsHandler> env_stats_handler;
   std::unique_ptr<FileRegistry> file_registry;
   std::vector<std::unique_ptr<rocksdb::Env>> envs;
 };

c-deps/libroach/include/libroach.h

@@ -304,8 +304,16 @@ typedef struct {
   int64_t pending_compaction_bytes_estimate;
 } DBStatsResult;
 
+// DBEnvStatsResult contains Env stats (filesystem layer).
+typedef struct {
+  // encryption status (CCL only).
+  // This is a serialized enginepbccl/stats.proto:EncryptionStatus
+  DBString encryption_status;
+} DBEnvStatsResult;
+
 DBStatus DBGetStats(DBEngine* db, DBStatsResult* stats);
 DBString DBGetCompactionStats(DBEngine* db);
+DBStatus DBGetEnvStats(DBEngine* db, DBEnvStatsResult* stats);
 
 typedef struct {
   int level;