From 69e58af0b93746246d38b2dc44853e4a07bf924d Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Fri, 11 Nov 2022 14:00:16 +0200 Subject: [PATCH] Add cockroachdb quota manager driver Signed-off-by: Juan Antonio Osorio --- .github/workflows/test.yaml | 2 +- cmd/trillian_log_server/main.go | 3 +- cmd/trillian_log_signer/main.go | 3 +- quota/crdbqm/common_test.go | 50 ++++++ quota/crdbqm/crdb_quota.go | 100 ++++++++++++ quota/crdbqm/crdb_quota_test.go | 260 ++++++++++++++++++++++++++++++++ quota/crdbqm/quota_provider.go | 50 ++++++ 7 files changed, 465 insertions(+), 3 deletions(-) create mode 100644 quota/crdbqm/common_test.go create mode 100644 quota/crdbqm/crdb_quota.go create mode 100644 quota/crdbqm/crdb_quota_test.go create mode 100644 quota/crdbqm/quota_provider.go diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index dd9d18807c..1de66c5a62 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -33,4 +33,4 @@ jobs: cache: true - name: Run tests - run: go test -v ./storage/crdb/... \ No newline at end of file + run: go test -v ./storage/crdb/... ./quota/crdbqm/... \ No newline at end of file diff --git a/cmd/trillian_log_server/main.go b/cmd/trillian_log_server/main.go index 2f9b81c03b..a2aa7fffee 100644 --- a/cmd/trillian_log_server/main.go +++ b/cmd/trillian_log_server/main.go @@ -50,7 +50,8 @@ import ( _ "github.com/google/trillian/storage/crdb" _ "github.com/google/trillian/storage/mysql" - // Load MySQL quota provider + // Load quota providers + _ "github.com/google/trillian/quota/crdbqm" _ "github.com/google/trillian/quota/mysqlqm" ) diff --git a/cmd/trillian_log_signer/main.go b/cmd/trillian_log_signer/main.go index 21feb5466c..d9de2eb16e 100644 --- a/cmd/trillian_log_signer/main.go +++ b/cmd/trillian_log_signer/main.go @@ -56,7 +56,8 @@ import ( _ "github.com/google/trillian/storage/crdb" _ "github.com/google/trillian/storage/mysql" - // Load MySQL quota provider + // Load quota providers + _ "github.com/google/trillian/quota/crdbqm" _ "github.com/google/trillian/quota/mysqlqm" ) diff --git a/quota/crdbqm/common_test.go b/quota/crdbqm/common_test.go new file mode 100644 index 0000000000..50eb0234a0 --- /dev/null +++ b/quota/crdbqm/common_test.go @@ -0,0 +1,50 @@ +// Copyright 2022 Trillian Authors +// +// Licensed 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. + +package crdbqm + +import ( + "os" + "testing" + + "github.com/cockroachdb/cockroach-go/v2/testserver" + "github.com/google/trillian/storage/testdb" + "k8s.io/klog/v2" +) + +func TestMain(m *testing.M) { + ts, err := testserver.NewTestServer() + if err != nil { + klog.Errorf("Failed to start test server: %v", err) + os.Exit(1) + } + defer ts.Stop() + + // reset the test server URL path. By default cockroach sets it + // to point to a default database, we don't want that. + dburl := ts.PGURL() + dburl.Path = "/" + + // Set the environment variable for the test server + os.Setenv(testdb.CockroachDBURIEnv, dburl.String()) + + if !testdb.CockroachDBAvailable() { + klog.Errorf("CockroachDB not available, skipping all CockroachDB storage tests") + return + } + + status := m.Run() + + os.Exit(status) +} diff --git a/quota/crdbqm/crdb_quota.go b/quota/crdbqm/crdb_quota.go new file mode 100644 index 0000000000..488b9903de --- /dev/null +++ b/quota/crdbqm/crdb_quota.go @@ -0,0 +1,100 @@ +// Copyright 2022 Trillian Authors +// +// Licensed 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. + +// Package crdbqm defines a CockroachDB-based quota.Manager implementation. +package crdbqm + +import ( + "context" + "database/sql" + "errors" + + "github.com/google/trillian/quota" +) + +const ( + // DefaultMaxUnsequenced is a suggested value for MaxUnsequencedRows. + // Note that this is a Global/Write quota suggestion, so it applies across trees. + DefaultMaxUnsequenced = 500000 // About 2h of non-stop signing at 70QPS. + + // TODO(jaosorior): Come up with a more optimal solution for CRDB, as this is + // linear and too costly. + countFromUnsequencedTable = "SELECT COUNT(*) FROM Unsequenced" +) + +// ErrTooManyUnsequencedRows is returned when tokens are requested but Unsequenced has grown +// beyond the configured limit. +var ErrTooManyUnsequencedRows = errors.New("too many unsequenced rows") + +// QuotaManager is a CockroachDB-based quota.Manager implementation. +// +// QuotaManager only implements Global/Write quotas, which is based on the number of Unsequenced +// rows (to be exact, tokens = MaxUnsequencedRows - actualUnsequencedRows). +// Other quotas are considered infinite. +type QuotaManager struct { + DB *sql.DB + MaxUnsequencedRows int +} + +// GetTokens implements quota.Manager.GetTokens. +// It doesn't actually reserve or retrieve tokens, instead it allows access based on the number of +// rows in the Unsequenced table. +func (m *QuotaManager) GetTokens(ctx context.Context, numTokens int, specs []quota.Spec) error { + for _, spec := range specs { + if spec.Group != quota.Global || spec.Kind != quota.Write { + continue + } + // Only allow global writes if Unsequenced is under the expected limit + count, err := m.countUnsequenced(ctx) + if err != nil { + return err + } + if count+numTokens > m.MaxUnsequencedRows { + return ErrTooManyUnsequencedRows + } + } + return nil +} + +// PutTokens implements quota.Manager.PutTokens. +// It's a noop for QuotaManager. +func (m *QuotaManager) PutTokens(ctx context.Context, numTokens int, specs []quota.Spec) error { + return nil +} + +// ResetQuota implements quota.Manager.ResetQuota. +// It's a noop for QuotaManager. +func (m *QuotaManager) ResetQuota(ctx context.Context, specs []quota.Spec) error { + return nil +} + +func (m *QuotaManager) countUnsequenced(ctx context.Context) (int, error) { + // table names are lowercase for some reason + rows, err := m.DB.QueryContext(ctx, countFromUnsequencedTable) + if err != nil { + return 0, err + } + defer rows.Close() + if !rows.Next() { + return 0, errors.New("cursor has no rows after quota limit determination query") + } + var count int + if err := rows.Scan(&count); err != nil { + return 0, err + } + if rows.Next() { + return 0, errors.New("too many rows returned from quota limit determination query") + } + return count, nil +} diff --git a/quota/crdbqm/crdb_quota_test.go b/quota/crdbqm/crdb_quota_test.go new file mode 100644 index 0000000000..feb5e44f96 --- /dev/null +++ b/quota/crdbqm/crdb_quota_test.go @@ -0,0 +1,260 @@ +// Copyright 2022 Trillian Authors +// +// Licensed 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. + +package crdbqm + +import ( + "context" + "crypto" + "database/sql" + "fmt" + "testing" + "time" + + "github.com/google/trillian" + "github.com/google/trillian/quota" + "github.com/google/trillian/storage" + "github.com/google/trillian/storage/crdb" + "github.com/google/trillian/storage/testdb" + "github.com/google/trillian/types" + + stestonly "github.com/google/trillian/storage/testonly" +) + +func TestQuotaManager_GetTokens(t *testing.T) { + testdb.SkipIfNoCockroachDB(t) + ctx := context.Background() + + db, done, err := testdb.NewTrillianDB(ctx, testdb.DriverCockroachDB) + if err != nil { + t.Fatalf("GetTestDB() returned err = %v", err) + } + defer done(ctx) + + tree, err := createTree(ctx, db) + if err != nil { + t.Fatalf("createTree() returned err = %v", err) + } + + tests := []struct { + desc string + unsequencedRows, maxUnsequencedRows, numTokens int + specs []quota.Spec + wantErr bool + }{ + { + desc: "globalWriteSingleToken", + unsequencedRows: 10, + maxUnsequencedRows: 20, + numTokens: 1, + specs: []quota.Spec{{Group: quota.Global, Kind: quota.Write}}, + }, + { + desc: "globalWriteMultiToken", + unsequencedRows: 10, + maxUnsequencedRows: 20, + numTokens: 5, + specs: []quota.Spec{{Group: quota.Global, Kind: quota.Write}}, + }, + { + desc: "globalWriteOverQuota1", + unsequencedRows: 20, + maxUnsequencedRows: 20, + numTokens: 1, + specs: []quota.Spec{{Group: quota.Global, Kind: quota.Write}}, + wantErr: true, + }, + { + desc: "globalWriteOverQuota2", + unsequencedRows: 15, + maxUnsequencedRows: 20, + numTokens: 10, + specs: []quota.Spec{{Group: quota.Global, Kind: quota.Write}}, + wantErr: true, + }, + { + desc: "unlimitedQuotas", + numTokens: 10, + specs: []quota.Spec{ + {Group: quota.User, Kind: quota.Read, User: "dylan"}, + {Group: quota.Tree, Kind: quota.Read, TreeID: tree.TreeId}, + {Group: quota.Global, Kind: quota.Read}, + {Group: quota.User, Kind: quota.Write, User: "dylan"}, + {Group: quota.Tree, Kind: quota.Write, TreeID: tree.TreeId}, + }, + }, + } + + for _, test := range tests { + t.Run(test.desc, func(t *testing.T) { + if err := setUnsequencedRows(ctx, db, tree, test.unsequencedRows); err != nil { + t.Errorf("setUnsequencedRows() returned err = %v", err) + return + } + + qm := &QuotaManager{DB: db, MaxUnsequencedRows: test.maxUnsequencedRows} + + err = qm.GetTokens(ctx, test.numTokens, test.specs) + if hasErr := err == ErrTooManyUnsequencedRows; hasErr != test.wantErr { + t.Errorf("%v: GetTokens() returned err = %q, wantErr = %v", test.desc, err, test.wantErr) + } + }) + } +} + +func TestQuotaManager_Noops(t *testing.T) { + testdb.SkipIfNoCockroachDB(t) + ctx := context.Background() + + db, done, err := testdb.NewTrillianDB(ctx, testdb.DriverCockroachDB) + if err != nil { + t.Fatalf("GetTestDB() returned err = %v", err) + } + defer done(ctx) + + qm := &QuotaManager{DB: db, MaxUnsequencedRows: 1000} + specs := allSpecs(ctx, qm, 10 /* treeID */) + + tests := []struct { + desc string + fn func() error + }{ + { + desc: "PutTokens", + fn: func() error { + return qm.PutTokens(ctx, 10 /* numTokens */, specs) + }, + }, + { + desc: "ResetQuota", + fn: func() error { + return qm.ResetQuota(ctx, specs) + }, + }, + } + for _, test := range tests { + if err := test.fn(); err != nil { + t.Errorf("%v: got err = %v", test.desc, err) + } + } +} + +func allSpecs(_ context.Context, _ quota.Manager, treeID int64) []quota.Spec { + return []quota.Spec{ + {Group: quota.User, Kind: quota.Read, User: "florence"}, + {Group: quota.Tree, Kind: quota.Read, TreeID: treeID}, + {Group: quota.Global, Kind: quota.Read}, + {Group: quota.User, Kind: quota.Write, User: "florence"}, + {Group: quota.Tree, Kind: quota.Write, TreeID: treeID}, + {Group: quota.Global, Kind: quota.Write}, + } +} + +func countUnsequenced(ctx context.Context, db *sql.DB) (int, error) { + var count int + if err := db.QueryRowContext(ctx, "SELECT COUNT(*) FROM Unsequenced").Scan(&count); err != nil { + return 0, err + } + return count, nil +} + +func createTree(ctx context.Context, db *sql.DB) (*trillian.Tree, error) { + var tree *trillian.Tree + + { + as := crdb.NewSQLAdminStorage(db) + err := as.ReadWriteTransaction(ctx, func(ctx context.Context, tx storage.AdminTX) error { + var err error + tree, err = tx.CreateTree(ctx, stestonly.LogTree) + return err + }) + if err != nil { + return nil, err + } + } + + { + ls := crdb.NewLogStorage(db, nil) + err := ls.ReadWriteTransaction(ctx, tree, func(ctx context.Context, tx storage.LogTreeTX) error { + logRoot, err := (&types.LogRootV1{RootHash: []byte{0}}).MarshalBinary() + if err != nil { + return err + } + slr := &trillian.SignedLogRoot{LogRoot: logRoot} + return tx.StoreSignedLogRoot(ctx, slr) + }) + if err != nil { + return nil, err + } + } + + return tree, nil +} + +func queueLeaves(ctx context.Context, db *sql.DB, tree *trillian.Tree, firstID, num int) error { + hasher := crypto.SHA256.New() + + leaves := []*trillian.LogLeaf{} + for i := 0; i < num; i++ { + value := []byte(fmt.Sprintf("leaf-%v", firstID+i)) + hasher.Reset() + if _, err := hasher.Write(value); err != nil { + return err + } + hash := hasher.Sum(nil) + leaves = append(leaves, &trillian.LogLeaf{ + MerkleLeafHash: hash, + LeafValue: value, + ExtraData: []byte("extra data"), + LeafIdentityHash: hash, + }) + } + + ls := crdb.NewLogStorage(db, nil) + _, err := ls.QueueLeaves(ctx, tree, leaves, time.Now()) + return err +} + +func setUnsequencedRows(ctx context.Context, db *sql.DB, tree *trillian.Tree, wantRows int) error { + count, err := countUnsequenced(ctx, db) + if err != nil { + return err + } + if count == wantRows { + return nil + } + + // Clear the tables and re-create leaves from scratch. It's easier than having to reason + // about duplicate entries. + if _, err := db.ExecContext(ctx, "DELETE FROM LeafData"); err != nil { + return err + } + if _, err := db.ExecContext(ctx, "DELETE FROM Unsequenced"); err != nil { + return err + } + if err := queueLeaves(ctx, db, tree, 0 /* firstID */, wantRows); err != nil { + return err + } + + // Sanity check the final count + count, err = countUnsequenced(ctx, db) + if err != nil { + return err + } + if count != wantRows { + return fmt.Errorf("got %v unsequenced rows, want = %v", count, wantRows) + } + + return nil +} diff --git a/quota/crdbqm/quota_provider.go b/quota/crdbqm/quota_provider.go new file mode 100644 index 0000000000..06b6fbbb30 --- /dev/null +++ b/quota/crdbqm/quota_provider.go @@ -0,0 +1,50 @@ +// Copyright 2022 Trillian Authors +// +// Licensed 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. + +package crdbqm + +import ( + "flag" + + "k8s.io/klog/v2" + + "github.com/google/trillian/quota" + "github.com/google/trillian/storage/crdb" +) + +// QuotaManagerName identifies the CockroachDB quota implementation. +const QuotaManagerName = "crdb" + +var maxUnsequencedRows = flag.Int("crdb_max_unsequenced_rows", DefaultMaxUnsequenced, + "Max number of unsequenced rows before rate limiting kicks in. Only effective for quota_system=crdb.") + +func init() { + if err := quota.RegisterProvider(QuotaManagerName, newCockroachDBQuotaManager); err != nil { + klog.Fatalf("Failed to register quota manager %v: %v", QuotaManagerName, err) + } +} + +func newCockroachDBQuotaManager() (quota.Manager, error) { + db, err := crdb.GetDatabase() + if err != nil { + return nil, err + } + qm := &QuotaManager{ + DB: db, + MaxUnsequencedRows: *maxUnsequencedRows, + } + + klog.Info("Using CockroachDB QuotaManager") + return qm, nil +}