Merge branch 'fix/784-correctness-fixes' into 'master'

fix: snapshot exemptions, observation stop hang, start timeout, token reload, panic guards (#784)

Closes #29 and #784

See merge request postgres-ai/database-lab!1199

Author: agneum

engine/cmd/cli/commands/clone/actions.go

@@ -132,7 +132,12 @@ func create(cliCtx *cli.Context) error {
 		cloneRequest.Snapshot = &types.SnapshotCloneFieldRequest{ID: cliCtx.String("snapshot-id")}
 	}
 
-	cloneRequest.ExtraConf = splitFlags(cliCtx.StringSlice("extra-config"))
+	extraConf, err := splitFlags(cliCtx.StringSlice("extra-config"))
+	if err != nil {
+		return fmt.Errorf("invalid --extra-config value: %w", err)
+	}
+
+	cloneRequest.ExtraConf = extraConf
 
 	var clone *models.Clone
 
@@ -373,10 +378,15 @@ func startObservation(cliCtx *cli.Context) error {
 		MaxDuration:         cliCtx.Uint64("max-duration"),
 	}
 
+	tags, err := splitFlags(cliCtx.StringSlice("tags"))
+	if err != nil {
+		return fmt.Errorf("invalid --tags value: %w", err)
+	}
+
 	start := types.StartObservationRequest{
 		CloneID: cloneID,
 		Config:  observationConfig,
-		Tags:    splitFlags(cliCtx.StringSlice("tags")),
+		Tags:    tags,
 		DBName:  cliCtx.String("db-name"),
 	}
 
@@ -569,19 +579,18 @@ func retrieveClonePort(cliCtx *cli.Context, wg *sync.WaitGroup, remoteHost *url.
 	return clone.DB.Port, nil
 }
 
-func splitFlags(flags []string) map[string]string {
-	const maxSplitParts = 2
-
+// splitFlags parses repeated key=value flag entries into a map; an entry without "=" is an error.
+func splitFlags(flags []string) (map[string]string, error) {
 	extraConfig := make(map[string]string, len(flags))
 
-	if len(flags) == 0 {
-		return extraConfig
-	}
-
 	for _, cfg := range flags {
-		parsed := strings.SplitN(cfg, "=", maxSplitParts)
-		extraConfig[parsed[0]] = parsed[1]
+		key, value, found := strings.Cut(cfg, "=")
+		if !found || key == "" {
+			return nil, fmt.Errorf("%q is not in the key=value form", cfg)
+		}
+
+		extraConfig[key] = value
 	}
 
-	return extraConfig
+	return extraConfig, nil
 }

engine/cmd/cli/commands/clone/actions_test.go

@@ -41,3 +41,36 @@ func TestUpgradeCommandIsRegistered(t *testing.T) {
 	assert.True(t, flagNames["async"])
 	assert.False(t, flagNames["target-version"], "the target follows from the instance configuration, not from a flag")
 }
+
+func TestSplitFlags(t *testing.T) {
+	testCases := []struct {
+		name    string
+		flags   []string
+		want    map[string]string
+		wantErr bool
+	}{
+		{name: "empty", flags: nil, want: map[string]string{}},
+		{name: "single pair", flags: []string{"shared_buffers=1GB"}, want: map[string]string{"shared_buffers": "1GB"}},
+		{name: "value keeps extra equals", flags: []string{"search_path=a=b"}, want: map[string]string{"search_path": "a=b"}},
+		{name: "empty value", flags: []string{"work_mem="}, want: map[string]string{"work_mem": ""}},
+		{name: "several pairs", flags: []string{"a=1", "b=2"}, want: map[string]string{"a": "1", "b": "2"}},
+		{name: "missing equals", flags: []string{"foo"}, wantErr: true},
+		{name: "missing key", flags: []string{"=bar"}, wantErr: true},
+		{name: "bad entry among good", flags: []string{"a=1", "foo"}, wantErr: true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := splitFlags(tc.flags)
+			if tc.wantErr {
+				assert.Error(t, err)
+				assert.Nil(t, got)
+
+				return
+			}
+
+			assert.NoError(t, err)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}

engine/internal/observer/observer_test.go

@@ -1,6 +1,7 @@
 package observer
 
 import (
+	"context"
 	"regexp"
 	"sync"
 	"testing"
@@ -9,7 +10,9 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/resources"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/client/dblabapi/types"
+	"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
 )
 
 func TestMaskingField(t *testing.T) {
@@ -219,3 +222,47 @@ func TestObservingClone_SetOverallError(t *testing.T) {
 	oc.SetOverallError(false)
 	assert.False(t, oc.session.state.OverallError)
 }
+
+func TestObservingClone_RunSessionErrorSignalsDone(t *testing.T) {
+	oc := NewObservingClone(types.Config{}, nil)
+
+	require.Error(t, oc.RunSession(), "a session that was never initialized cannot run")
+
+	select {
+	case <-oc.done:
+	case <-time.After(time.Second):
+		t.Fatal("done must be signalled when RunSession exits with an error")
+	}
+}
+
+func TestObservingClone_StopExpiredContext(t *testing.T) {
+	oc := NewObservingClone(types.Config{}, nil)
+	oc.session = &Session{SessionID: 42}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err := oc.Stop(ctx)
+
+	require.Error(t, err, "Stop must not block when RunSession never started")
+	assert.ErrorIs(t, err, context.Canceled)
+}
+
+func TestObservingClone_StopUninitializedSession(t *testing.T) {
+	oc := NewObservingClone(types.Config{}, nil)
+
+	require.Error(t, oc.Stop(context.Background()))
+}
+
+func TestObservingClone_StopTwiceAfterRunSessionExit(t *testing.T) {
+	oc := NewObservingClone(types.Config{}, nil)
+	oc.pool = &resources.Pool{Name: "pool", MountDir: t.TempDir()}
+	oc.session = &Session{SessionID: 7, Config: types.Config{MaxDuration: 60}, Result: &models.ObservationResult{}}
+
+	close(oc.done)
+
+	require.NotPanics(t, func() {
+		require.NoError(t, oc.Stop(context.Background()))
+		require.NoError(t, oc.Stop(context.Background()))
+	})
+}

engine/internal/observer/observing_clone.go

@@ -97,7 +97,7 @@ func NewObservingClone(config types.Config, sudb *pgx.Conn) *ObservingClone {
 		config:          config,
 		ctx:             ctx,
 		cancel:          cancel,
-		done:            make(chan struct{}, 1),
+		done:            make(chan struct{}),
 		csvFields:       csvFields,
 		registryMu:      &sync.Mutex{},
 		sessionRegistry: make(map[uint64]struct{}),
@@ -202,8 +202,11 @@ func (c *ObservingClone) Init(clone *models.Clone, sessionID uint64, startedAt t
 	return nil
 }
 
-// RunSession runs observing session.
+// RunSession runs observing session. It closes the done channel on every return path so that
+// Stop never waits for a session that has already ended.
 func (c *ObservingClone) RunSession() error {
+	defer close(c.done)
+
 	if c.session == nil || c.db.IsClosed() {
 		return errors.New("failed to run session because it has not been initialized")
 	}
@@ -259,8 +262,6 @@ func (c *ObservingClone) RunSession() error {
 				log.Err("failed to store artifacts: ", err)
 			}
 
-			c.done <- struct{}{}
-
 			return nil
 		}
 
@@ -436,16 +437,20 @@ where table_name = 'postgres_log'`)
 }
 
 // Stop stops an observation session.
-func (c *ObservingClone) Stop() error {
+func (c *ObservingClone) Stop(ctx context.Context) error {
+	if c.session == nil {
+		return errors.New("failed to summarize session because it has not been initialized")
+	}
+
 	log.Msg(fmt.Sprintf("Observation session %v is stopping...", c.session.SessionID))
 
 	c.cancel()
 
-	// Waiting for the observation process stops.
-	<-c.done
-
-	if c.session == nil {
-		return errors.New("failed to summarize session because it has not been initialized")
+	// Wait until the observation process ends or the caller gives up.
+	select {
+	case <-c.done:
+	case <-ctx.Done():
+		return fmt.Errorf("observation session %v has not stopped: %w", c.session.SessionID, ctx.Err())
 	}
 
 	c.summarize()

engine/internal/provision/databases/postgres/postgres.go

@@ -119,7 +119,7 @@ func Start(r runners.Runner, c *resources.AppConfig) error {
 				log.Err(runnerErr)
 			}
 
-			return errors.Wrap(err, "postgres start timeout")
+			return startTimeoutError(out, err)
 		}
 
 		time.Sleep(checkPostgresStatusPeriod * time.Millisecond)
@@ -128,6 +128,16 @@ func Start(r runners.Runner, c *resources.AppConfig) error {
 	return nil
 }
 
+// startTimeoutError describes why Postgres was not ready before the deadline: recoveryState is
+// the last pg_is_in_recovery() result and queryErr the error of the last status query, if any.
+func startTimeoutError(recoveryState string, queryErr error) error {
+	if queryErr != nil {
+		return fmt.Errorf("postgres start timeout: last status check failed: %w", queryErr)
+	}
+
+	return fmt.Errorf("postgres start timeout: instance is still in recovery (pg_is_in_recovery: %q)", recoveryState)
+}
+
 func collectDiagnostics(c *resources.AppConfig) {
 	dockerClient, err := client.New(client.FromEnv)
 	if err != nil {

engine/internal/provision/databases/postgres/postgres_test.go

@@ -170,3 +170,32 @@ func TestRemoveContainers(t *testing.T) {
 		assert.Equal(t, tc.err, errors.Cause(err))
 	}
 }
+
+func TestStartTimeoutError(t *testing.T) {
+	queryErr := errors.New("connection refused")
+
+	testCases := []struct {
+		name          string
+		recoveryState string
+		queryErr      error
+		wantContains  string
+	}{
+		{name: "still in recovery", recoveryState: "t", queryErr: nil, wantContains: `still in recovery (pg_is_in_recovery: "t")`},
+		{name: "no status at all", recoveryState: "", queryErr: nil, wantContains: `pg_is_in_recovery: ""`},
+		{name: "status check failed", recoveryState: "", queryErr: queryErr, wantContains: "last status check failed: connection refused"},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err := startTimeoutError(tc.recoveryState, tc.queryErr)
+
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "postgres start timeout")
+			assert.Contains(t, err.Error(), tc.wantContains)
+
+			if tc.queryErr != nil {
+				assert.ErrorIs(t, err, tc.queryErr)
+			}
+		})
+	}
+}

engine/internal/retrieval/retrieval.go

@@ -368,9 +368,7 @@ func (r *Retrieval) run(ctx context.Context, fsm pool.FSManager) (err error) {
 		r.State.cleanAlerts()
 	}
 
-	var existsErr *thinclones.SnapshotExistsError
-
-	if err := r.SnapshotData(ctx, poolName); err != nil && (err != errNoJobs || !errors.As(err, &existsErr)) {
+	if err := r.SnapshotData(ctx, poolName); err != nil && !isSnapshotExempt(err) {
 		return err
 	}
 
@@ -390,6 +388,14 @@ func (r *Retrieval) run(ctx context.Context, fsm pool.FSManager) (err error) {
 	return nil
 }
 
+// isSnapshotExempt reports whether a SnapshotData error must not abort the run: having no
+// snapshot jobs or an already existing snapshot still leaves the pool ready to be activated.
+func isSnapshotExempt(err error) bool {
+	var existsErr *thinclones.SnapshotExistsError
+
+	return errors.Is(err, errNoJobs) || errors.As(err, &existsErr)
+}
+
 // RefreshData runs a group of data refresh jobs.
 func (r *Retrieval) RefreshData(ctx context.Context, poolName string) error {
 	fsm, err := r.poolManager.GetFSManager(poolName)

engine/internal/retrieval/retrieval_test.go

@@ -2,13 +2,16 @@ package retrieval
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"os"
 	"path"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"gitlab.com/postgres-ai/database-lab/v3/internal/provision/thinclones"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/config"
 	"gitlab.com/postgres-ai/database-lab/v3/internal/retrieval/engine/postgres/logical"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
@@ -306,3 +309,23 @@ func TestSkipRefreshingError(t *testing.T) {
 		assert.EqualError(t, err, "some error")
 	})
 }
+
+func TestIsSnapshotExempt(t *testing.T) {
+	testCases := []struct {
+		name   string
+		err    error
+		exempt bool
+	}{
+		{name: "no jobs", err: errNoJobs, exempt: true},
+		{name: "wrapped no jobs", err: fmt.Errorf("snapshot: %w", errNoJobs), exempt: true},
+		{name: "snapshot exists", err: thinclones.NewSnapshotExistsError("snap"), exempt: true},
+		{name: "wrapped snapshot exists", err: fmt.Errorf("snapshot: %w", thinclones.NewSnapshotExistsError("snap")), exempt: true},
+		{name: "other error", err: errors.New("zfs failed"), exempt: false},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.exempt, isSnapshotExempt(tc.err))
+		})
+	}
+}

engine/internal/runci/server.go

@@ -54,7 +54,7 @@ func NewServer(cfg *Config, dle *dblabapi.Client, platform *platform.Service, co
 func (s *Server) Run() error {
 	r := mux.NewRouter().StrictSlash(true)
 
-	authMW := mw.NewAuth(s.config.App.VerificationToken, s.platform)
+	authMW := mw.NewAuth(mw.StaticToken(s.config.App.VerificationToken), s.platform)
 
 	r.HandleFunc("/migration/run", authMW.Authorized(s.runMigration)).Methods(http.MethodPost)
 	r.HandleFunc("/artifact/download", authMW.Authorized(s.downloadArtifact)).Methods(http.MethodGet)
@@ -63,7 +63,7 @@ func (s *Server) Run() error {
 
 	addr := fmt.Sprintf("%s:%d", s.config.App.Host, s.config.App.Port)
 
-	s.httpServer = &http.Server{Addr: addr, Handler: mw.Logging(r)}
+	s.httpServer = &http.Server{Addr: addr, Handler: mw.Logging(mw.Recover(r))}
 
 	log.Msg(fmt.Sprintf("Server started listening on %s...", addr))