From df991698ecfec07a5afb8009e96331bc197f910c Mon Sep 17 00:00:00 2001 From: Aaron Son Date: Sun, 23 Feb 2025 11:46:12 -0800 Subject: [PATCH] go: sqle/auto_gc: Add some simple tests for controller and hook. --- go/cmd/dolt/commands/engine/sqlengine.go | 1 + go/libraries/doltcore/sqle/auto_gc.go | 13 ++- go/libraries/doltcore/sqle/auto_gc_test.go | 117 +++++++++++++++++++++ 3 files changed, 127 insertions(+), 4 deletions(-) create mode 100644 go/libraries/doltcore/sqle/auto_gc_test.go diff --git a/go/cmd/dolt/commands/engine/sqlengine.go b/go/cmd/dolt/commands/engine/sqlengine.go index fe1769c0ca..9c43cab2b9 100644 --- a/go/cmd/dolt/commands/engine/sqlengine.go +++ b/go/cmd/dolt/commands/engine/sqlengine.go @@ -211,6 +211,7 @@ func NewSqlEngine( } config.AutoGCController.ApplyCommitHooks(ctx, mrEnv, dbs...) pro.InitDatabaseHooks = append(pro.InitDatabaseHooks, config.AutoGCController.InitDatabaseHook()) + pro.DropDatabaseHooks = append(pro.DropDatabaseHooks, config.AutoGCController.DropDatabaseHook()) // XXX: We force session aware safepoint controller if auto_gc is on. dprocedures.UseSessionAwareSafepointController = true } diff --git a/go/libraries/doltcore/sqle/auto_gc.go b/go/libraries/doltcore/sqle/auto_gc.go index 5ca099bd49..ac33a42a17 100644 --- a/go/libraries/doltcore/sqle/auto_gc.go +++ b/go/libraries/doltcore/sqle/auto_gc.go @@ -83,7 +83,12 @@ func (c *AutoGCController) RunBackgroundThread(threads *sql.BackgroundThreads, c if err != nil { return err } - // TODO: Start bg threads for all existing commit hooks. + for _, hook := range c.hooks { + err = hook.run(threads) + if err != nil { + return err + } + } return nil } @@ -157,7 +162,7 @@ func (c *AutoGCController) doWork(ctx context.Context, work autoGCWork, ctxF fun c.lgr.Infof("sqle/auto_gc: Successfully completed auto GC of database %s in %v", work.name, time.Since(start)) } -func (c *AutoGCController) newCommitHook(name string, db *doltdb.DoltDB) doltdb.CommitHook { +func (c *AutoGCController) newCommitHook(name string, db *doltdb.DoltDB) *autoGCCommitHook { c.mu.Lock() defer c.mu.Unlock() closed := make(chan struct{}) @@ -231,7 +236,7 @@ func (c *AutoGCController) ApplyCommitHooks(ctx context.Context, mrEnv *env.Mult } func (c *AutoGCController) DropDatabaseHook() DropDatabaseHook { - return func(ctx *sql.Context, name string) { + return func(_ *sql.Context, name string) { c.mu.Lock() defer c.mu.Unlock() hook := c.hooks[name] @@ -243,7 +248,7 @@ func (c *AutoGCController) DropDatabaseHook() DropDatabaseHook { } func (c *AutoGCController) InitDatabaseHook() InitDatabaseHook { - return func(ctx *sql.Context, pro *DoltDatabaseProvider, name string, env *env.DoltEnv, db dsess.SqlDatabase) error { + return func(ctx *sql.Context, _ *DoltDatabaseProvider, name string, env *env.DoltEnv, _ dsess.SqlDatabase) error { ddb := env.DoltDB(ctx) ddb.PrependCommitHooks(ctx, c.newCommitHook(name, ddb)) return nil diff --git a/go/libraries/doltcore/sqle/auto_gc_test.go b/go/libraries/doltcore/sqle/auto_gc_test.go new file mode 100644 index 0000000000..80387f57b2 --- /dev/null +++ b/go/libraries/doltcore/sqle/auto_gc_test.go @@ -0,0 +1,117 @@ +// Copyright 2025 Dolthub, Inc. +// +// 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 sqle + +import ( + "bytes" + "context" + "sync" + "testing" + "time" + + "github.com/sirupsen/logrus" + "github.com/stretchr/testify/require" + + "github.com/dolthub/dolt/go/store/datas" + "github.com/dolthub/go-mysql-server/sql" +) + +func TestAutoGCController(t *testing.T) { + NewLogger := func() *logrus.Logger { + res := logrus.New() + res.SetOutput(new(bytes.Buffer)) + return res + } + CtxFactory := func(ctx context.Context) (*sql.Context, error) { + return sql.NewContext(ctx, sql.WithSession(sql.NewBaseSession())), nil + } + t.Run("Hook", func(t *testing.T) { + t.Run("NeverStarted", func(t *testing.T) { + controller := NewAutoGCController(NewLogger()) + hook := controller.newCommitHook("some_database", nil) + hook.stop() + }) + t.Run("StartedBeforeNewHook", func(t *testing.T) { + controller := NewAutoGCController(NewLogger()) + bg := sql.NewBackgroundThreads() + defer bg.Shutdown() + err := controller.RunBackgroundThread(bg, CtxFactory) + require.NoError(t, err) + ctx := context.Background() + dEnv := CreateTestEnvWithName("some_database") + hook := controller.newCommitHook("some_database", dEnv.DoltDB(ctx)) + hook.Execute(ctx, datas.Dataset{}, nil) + hook.stop() + }) + t.Run("StartedAfterNewHook", func(t *testing.T) { + controller := NewAutoGCController(NewLogger()) + bg := sql.NewBackgroundThreads() + defer bg.Shutdown() + ctx := context.Background() + dEnv := CreateTestEnvWithName("some_database") + hook := controller.newCommitHook("some_database", dEnv.DoltDB(ctx)) + err := controller.RunBackgroundThread(bg, CtxFactory) + require.NoError(t, err) + hook.Execute(ctx, datas.Dataset{}, nil) + hook.stop() + }) + t.Run("ExecuteOnCanceledCtx", func(t *testing.T) { + controller := NewAutoGCController(NewLogger()) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + dEnv := CreateTestEnvWithName("some_database") + hook := controller.newCommitHook("some_database", dEnv.DoltDB(ctx)) + _, err := hook.Execute(ctx, datas.Dataset{}, nil) + require.ErrorIs(t, err, context.Canceled) + }) + }) + t.Run("gcBgThread", func(t *testing.T) { + controller := NewAutoGCController(NewLogger()) + ctx, cancel := context.WithCancel(context.Background()) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + controller.gcBgThread(ctx) + }() + time.Sleep(50 * time.Millisecond) + cancel() + wg.Wait() + }) + t.Run("DatabaseProviderHooks", func(t *testing.T) { + t.Run("Unstarted", func(t *testing.T) { + controller := NewAutoGCController(NewLogger()) + ctx, err := CtxFactory(context.Background()) + require.NoError(t, err) + dEnv := CreateTestEnvWithName("some_database") + err = controller.InitDatabaseHook()(ctx, nil, "some_database", dEnv, nil) + require.NoError(t, err) + controller.DropDatabaseHook()(nil, "some_database") + }) + t.Run("Started", func(t *testing.T) { + controller := NewAutoGCController(NewLogger()) + bg := sql.NewBackgroundThreads() + defer bg.Shutdown() + err := controller.RunBackgroundThread(bg, CtxFactory) + require.NoError(t, err) + ctx, err := CtxFactory(context.Background()) + require.NoError(t, err) + dEnv := CreateTestEnvWithName("some_database") + err = controller.InitDatabaseHook()(ctx, nil, "some_database", dEnv, nil) + require.NoError(t, err) + controller.DropDatabaseHook()(nil, "some_database") + }) + }) +}