HH-807: harden runtime coordination (#21)

internal/hub/store.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"net/url"
 	"regexp"
+	"sort"
 	"strings"
 	"time"
 	"unicode/utf8"
@@ -62,6 +63,8 @@ var (
 
 func ValidImageVersion(version string) bool { return imageVersionPattern.MatchString(version) }
 
+func ValidNetworkExitID(id string) bool { return exitIDPattern.MatchString(id) }
+
 var (
 	aliasPattern        = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,31}$`)
 	gatewayNamePattern  = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$`)
@@ -71,8 +74,9 @@ var (
 )
 
 type Store struct {
-	db     *sql.DB
-	notify taskstate.Notifier
+	db            *sql.DB
+	lockAdmission chan struct{}
+	notify        taskstate.Notifier
 }
 
 // Gateway 是平台注册的 docker-gateway 实例；Token 由平台生成，明文存储供页面复制（开发阶段约定）。
@@ -116,7 +120,9 @@ func Open(ctx context.Context, databaseURL string) (*Store, error) {
 		db.Close()
 		return nil, errors.New("connect to hub database")
 	}
-	store := &Store{db: db}
+	// LockResources keeps one connection until the lifecycle operation finishes.
+	// Admit at most half the pool so those operations can still open nested DB calls.
+	store := &Store{db: db, lockAdmission: make(chan struct{}, 5)}
 	if err := store.migrate(ctx); err != nil {
 		db.Close()
 		return nil, err
@@ -137,6 +143,58 @@ func (s *Store) notifyTransitions(transitions []taskstate.Transition) {
 	}
 }
 
+// LockResources serializes lifecycle state across control-plane replicas. The
+// transaction carries no data changes; rolling it back only releases the locks.
+func (s *Store) LockResources(ctx context.Context, aliases, exitIDs, imageVersions []string) (func(), error) {
+	for _, alias := range aliases {
+		if !aliasPattern.MatchString(alias) {
+			return nil, ErrInvalid
+		}
+	}
+	for _, id := range exitIDs {
+		if !exitIDPattern.MatchString(id) {
+			return nil, ErrInvalid
+		}
+	}
+	for _, version := range imageVersions {
+		if !imageVersionPattern.MatchString(version) {
+			return nil, ErrInvalid
+		}
+	}
+	if len(aliases)+len(exitIDs)+len(imageVersions) == 0 {
+		return func() {}, nil
+	}
+	select {
+	case s.lockAdmission <- struct{}{}:
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+	tx, err := s.db.BeginTx(ctx, nil)
+	if err != nil {
+		<-s.lockAdmission
+		return nil, errors.New("begin resource lock")
+	}
+	resources := []struct {
+		namespace int
+		keys      []string
+	}{{1542738013, aliases}, {1542738015, exitIDs}, {1542738014, imageVersions}}
+	for _, resource := range resources {
+		keys := append([]string(nil), resource.keys...)
+		sort.Strings(keys)
+		for _, key := range keys {
+			if _, err := tx.ExecContext(ctx, `SELECT pg_advisory_xact_lock($1, hashtext($2))`, resource.namespace, key); err != nil {
+				_ = tx.Rollback()
+				<-s.lockAdmission
+				return nil, errors.New("lock resource")
+			}
+		}
+	}
+	return func() {
+		_ = tx.Rollback()
+		<-s.lockAdmission
+	}, nil
+}
+
 func (s *Store) migrate(ctx context.Context) error {
 	tx, err := s.db.BeginTx(ctx, nil)
 	if err != nil {

internal/hub/store_test.go

@@ -4,16 +4,137 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"os"
 	"reflect"
 	"slices"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"git.ipao.vip/rogee/creator-hub/internal/taskstate"
 )
 
+func TestEnvironmentLocksCoordinateAcrossStoreInstances(t *testing.T) {
+	databaseURL := os.Getenv("CREATORHUB_POSTGRES_TEST_URL")
+	if databaseURL == "" {
+		t.Skip("set CREATORHUB_POSTGRES_TEST_URL to run PostgreSQL integration coverage")
+	}
+	ctx := context.Background()
+	testURL := isolatedDatabaseURL(t, databaseURL)
+	first := openFullyMigratedHub(t, ctx, testURL)
+	t.Cleanup(func() { _ = first.Close() })
+	second, err := Open(ctx, testURL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = second.Close() })
+
+	unlockFirst, err := first.LockResources(ctx, []string{"account-a"}, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	firstReleased := false
+	defer func() {
+		if !firstReleased {
+			unlockFirst()
+		}
+	}()
+
+	differentAlias, err := second.LockResources(ctx, []string{"account-b"}, nil, nil)
+	if err != nil {
+		t.Fatalf("different aliases must not share a lock: %v", err)
+	}
+	differentAlias()
+
+	acquired := make(chan func(), 1)
+	errors := make(chan error, 1)
+	started := make(chan struct{})
+	go func() {
+		close(started)
+		unlock, lockErr := second.LockResources(ctx, []string{"account-a"}, nil, nil)
+		if lockErr != nil {
+			errors <- lockErr
+			return
+		}
+		acquired <- unlock
+	}()
+	<-started
+	select {
+	case unlock := <-acquired:
+		unlock()
+		t.Fatal("same alias lock did not block across Store instances")
+	case err := <-errors:
+		t.Fatal(err)
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	unlockFirst()
+	firstReleased = true
+	select {
+	case unlock := <-acquired:
+		unlock()
+	case err := <-errors:
+		t.Fatal(err)
+	case <-time.After(time.Second):
+		t.Fatal("same alias lock was not released")
+	}
+}
+
+func TestResourceLocksReserveConnectionsForLifecycleQueries(t *testing.T) {
+	databaseURL := os.Getenv("CREATORHUB_POSTGRES_TEST_URL")
+	if databaseURL == "" {
+		t.Skip("set CREATORHUB_POSTGRES_TEST_URL to run PostgreSQL integration coverage")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	store := openFullyMigratedHub(t, ctx, isolatedDatabaseURL(t, databaseURL))
+	t.Cleanup(func() { _ = store.Close() })
+
+	const workers = 10
+	var acquired atomic.Int32
+	startQueries := make(chan struct{})
+	done := make(chan error, workers)
+	for worker := 0; worker < workers; worker++ {
+		go func(worker int) {
+			unlock, err := store.LockResources(ctx, []string{fmt.Sprintf("account-%d", worker)}, nil, nil)
+			if err != nil {
+				done <- err
+				return
+			}
+			defer unlock()
+			acquired.Add(1)
+			<-startQueries
+			_, err = store.ListEnvs(ctx)
+			done <- err
+		}(worker)
+	}
+	deadline := time.NewTimer(100 * time.Millisecond)
+	ticker := time.NewTicker(time.Millisecond)
+	for acquired.Load() < workers {
+		select {
+		case <-ticker.C:
+		case <-deadline.C:
+			goto release
+		}
+	}
+release:
+	ticker.Stop()
+	if !deadline.Stop() {
+		select {
+		case <-deadline.C:
+		default:
+		}
+	}
+	close(startQueries)
+	for worker := 0; worker < workers; worker++ {
+		if err := <-done; err != nil {
+			t.Fatalf("locked lifecycle query %d did not complete: %v", worker, err)
+		}
+	}
+}
+
 func TestFingerprintArgsFollowUpstreamCommandLineContract(t *testing.T) {
 	full := Fingerprint{
 		Seed: 2024, Platform: "windows", PlatformVersion: "11.0.0",