From 0b806b47883e08087d47e3d84b7986ea2191185c Mon Sep 17 00:00:00 2001
From: Andy Yang <ayang@cockroachlabs.com>
Date: Wed, 9 Jun 2021 14:41:08 -0400
Subject: [PATCH] storageccl: don't update file registry when creating
 unencrypted files

This patch changes PebbleFileRegistry to not store an entry
for an unencrypted file. This should lead to a performance
improvement since the file registry is rewritten every time
it is updated.

Release note: None
---
 .../storageccl/engineccl/encrypted_fs_test.go | 99 +++++++++++++++++++
 pkg/storage/pebble_file_registry.go           |  6 ++
 2 files changed, 105 insertions(+)

diff --git a/pkg/ccl/storageccl/engineccl/encrypted_fs_test.go b/pkg/ccl/storageccl/engineccl/encrypted_fs_test.go
index 0437ee84bd48..93814a388c71 100644
--- a/pkg/ccl/storageccl/engineccl/encrypted_fs_test.go
+++ b/pkg/ccl/storageccl/engineccl/encrypted_fs_test.go
@@ -23,6 +23,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
 	"github.com/cockroachdb/pebble"
@@ -280,3 +281,101 @@ func TestPebbleEncryption(t *testing.T) {
 
 	db.Close()
 }
+
+func TestPebbleEncryption2(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	memFS := vfs.NewMem()
+	firstKeyFile128 := "111111111111111111111111111111111234567890123456"
+	secondKeyFile128 := "111111111111111111111111111111198765432198765432"
+	writeToFile(t, memFS, "16v1.key", []byte(firstKeyFile128))
+	writeToFile(t, memFS, "16v2.key", []byte(secondKeyFile128))
+
+	keys := make(map[string]bool)
+	validateKeys := func(reader storage.Reader) bool {
+		keysCopy := make(map[string]bool)
+		for k, v := range keys {
+			keysCopy[k] = v
+		}
+
+		foundUnknown := false
+		kvFunc := func(kv roachpb.KeyValue) error {
+			key := kv.Key
+			val := kv.Value
+			expected := keysCopy[string(key)]
+			if !expected || len(val.RawBytes) == 0 {
+				foundUnknown = true
+				return nil
+			}
+			delete(keysCopy, string(key))
+			return nil
+		}
+
+		_, err := storage.MVCCIterate(
+			context.Background(),
+			reader,
+			nil,
+			roachpb.KeyMax,
+			hlc.Timestamp{},
+			storage.MVCCScanOptions{},
+			kvFunc,
+		)
+		require.NoError(t, err)
+		return len(keysCopy) == 0 && !foundUnknown
+	}
+
+	addKeyAndValidate := func(
+		key string, val string, encKeyFile string, oldEncFileKey string,
+	) {
+		encOptions := baseccl.EncryptionOptions{
+			KeySource: baseccl.EncryptionKeySource_KeyFiles,
+			KeyFiles: &baseccl.EncryptionKeyFiles{
+				CurrentKey: encKeyFile,
+				OldKey:     oldEncFileKey,
+			},
+			DataKeyRotationPeriod: 1000,
+		}
+		encOptionsBytes, err := protoutil.Marshal(&encOptions)
+		require.NoError(t, err)
+
+		opts := storage.DefaultPebbleOptions()
+		opts.FS = memFS
+		opts.Cache = pebble.NewCache(1 << 20)
+		defer opts.Cache.Unref()
+
+		db, err := storage.NewPebble(
+			context.Background(),
+			storage.PebbleConfig{
+				StorageConfig: base.StorageConfig{
+					Attrs:             roachpb.Attributes{},
+					MaxSize:           512 << 20,
+					UseFileRegistry:   true,
+					EncryptionOptions: encOptionsBytes,
+				},
+				Opts: opts,
+			})
+		require.NoError(t, err)
+		require.True(t, validateKeys(db))
+
+		keys[key] = true
+		err = storage.MVCCPut(
+			context.Background(),
+			db,
+			nil, /* ms */
+			roachpb.Key(key),
+			hlc.Timestamp{},
+			roachpb.MakeValueFromBytes([]byte(val)),
+			nil, /* txn */
+		)
+		require.NoError(t, err)
+		require.NoError(t, db.Flush())
+		require.NoError(t, db.Compact())
+		require.True(t, validateKeys(db))
+		db.Close()
+	}
+
+	addKeyAndValidate("a", "a", "16v1.key", "plain")
+	addKeyAndValidate("b", "b", "plain", "16v1.key")
+	addKeyAndValidate("c", "c", "16v2.key", "plain")
+	addKeyAndValidate("d", "d", "plain", "16v2.key")
+}
diff --git a/pkg/storage/pebble_file_registry.go b/pkg/storage/pebble_file_registry.go
index 484882cce475..1ca2d43d315c 100644
--- a/pkg/storage/pebble_file_registry.go
+++ b/pkg/storage/pebble_file_registry.go
@@ -102,6 +102,12 @@ func (r *PebbleFileRegistry) GetFileEntry(filename string) *enginepb.FileEntry {
 
 // SetFileEntry sets filename => entry in the registry map and persists the registry.
 func (r *PebbleFileRegistry) SetFileEntry(filename string, entry *enginepb.FileEntry) error {
+	// We choose not to store an entry for unencrypted files since the absence of
+	// a file in the file registry implies that it is unencrypted.
+	if entry != nil && entry.EnvType == enginepb.EnvType_Plaintext {
+		return r.MaybeDeleteEntry(filename)
+	}
+
 	filename = r.tryMakeRelativePath(filename)
 	newProto := &enginepb.FileRegistry{}