Reduce catalog contention for RecordMetricsCritical

flow/activities/flowable_core.go

@@ -559,12 +559,9 @@ func replicateQRepPartition[TRead any, TWrite QRepStreamCloser, TSync connectors
 
 	if rowsSynced > 0 {
 		logger.Info(fmt.Sprintf("pushed %d records", rowsSynced))
-		if err := monitoring.UpdateRowsSyncedForPartition(ctx, a.CatalogPool, rowsSynced, runUUID, partition); err != nil {
-			return err
-		}
 	}
 
-	return monitoring.UpdateEndTimeForPartition(ctx, a.CatalogPool, runUUID, partition)
+	return monitoring.UpdateEndTimeAndRowsSyncedForPartition(ctx, a.CatalogPool, rowsSynced, runUUID, partition)
 }
 
 // replicateXminPartition replicates a XminPartition from the source to the destination.
@@ -675,15 +672,10 @@ func replicateXminPartition[TRead any, TWrite QRepStreamCloser, TSync connectors
 	}
 
 	if rowsSynced > 0 {
-		err := monitoring.UpdateRowsSyncedForPartition(ctx, a.CatalogPool, rowsSynced, runUUID, partition)
-		if err != nil {
-			return 0, err
-		}
-
 		logger.Info(fmt.Sprintf("pushed %d records", rowsSynced))
 	}
 
-	if err := monitoring.UpdateEndTimeForPartition(ctx, a.CatalogPool, runUUID, partition); err != nil {
+	if err := monitoring.UpdateEndTimeAndRowsSyncedForPartition(ctx, a.CatalogPool, rowsSynced, runUUID, partition); err != nil {
 		return 0, err
 	}
 

flow/connectors/utils/monitoring/monitoring.go

@@ -458,26 +458,16 @@ func UpdatePullEndTimeAndRowsForPartition(ctx context.Context, pool shared.Catal
 	return nil
 }
 
-func UpdateEndTimeForPartition(ctx context.Context, pool shared.CatalogPool, runUUID string,
+func UpdateEndTimeAndRowsSyncedForPartition(ctx context.Context, pool shared.CatalogPool, rowsSynced int64, runUUID string,
 	partition *protos.QRepPartition,
 ) error {
 	if _, err := pool.Exec(ctx,
-		`UPDATE peerdb_stats.qrep_partitions SET end_time=$1 WHERE run_uuid=$2 AND partition_uuid=$3`,
-		time.Now(), runUUID, partition.PartitionId,
+		`UPDATE peerdb_stats.qrep_partitions
+		SET end_time=$1, rows_synced=CASE WHEN $2::bigint > 0 THEN $2 ELSE rows_synced END
+		WHERE run_uuid=$3 AND partition_uuid=$4`,
+		time.Now(), rowsSynced, runUUID, partition.PartitionId,
 	); err != nil {
-		return fmt.Errorf("error while updating qrep partition in qrep_partitions: %w", err)
-	}
-	return nil
-}
-
-func UpdateRowsSyncedForPartition(ctx context.Context, pool shared.CatalogPool, rowsSynced int64, runUUID string,
-	partition *protos.QRepPartition,
-) error {
-	if _, err := pool.Exec(ctx,
-		`UPDATE peerdb_stats.qrep_partitions SET rows_synced=$1 WHERE run_uuid=$2 AND partition_uuid=$3`,
-		rowsSynced, runUUID, partition.PartitionId,
-	); err != nil {
-		return fmt.Errorf("error while updating rows_synced in qrep_partitions: %w", err)
+		return fmt.Errorf("error while completing qrep partition in qrep_partitions: %w", err)
 	}
 	return nil
 }

flow/internal/catalog.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/jackc/pgx/v5/pgxpool"
@@ -15,35 +16,39 @@ import (
 
 var (
 	poolMutex = &sync.Mutex{}
-	pool      *pgxpool.Pool
+	pool      atomic.Pointer[pgxpool.Pool]
 )
 
 func GetCatalogConnectionPoolFromEnv(ctx context.Context) (shared.CatalogPool, error) {
-	poolMutex.Lock()
-	defer poolMutex.Unlock()
-	if pool == nil {
-		var err error
-		catalogConnectionString := GetCatalogConnectionStringFromEnv(ctx)
-		config, err := pgxpool.ParseConfig(catalogConnectionString)
-		if err != nil {
-			return shared.CatalogPool{},
-				exceptions.NewCatalogError(fmt.Errorf("unable to parse catalog connection string: %w", err))
-		}
-		config.MaxConns = 3
-		config.MaxConnIdleTime = 90 * time.Second
-		pool, err = pgxpool.NewWithConfig(ctx, config)
-		if err != nil {
-			return shared.CatalogPool{Pool: pool},
-				exceptions.NewCatalogError(fmt.Errorf("unable to establish connection with catalog: %w", err))
-		}
-	}
+	if pool.Load() == nil {
+		poolMutex.Lock()
+		defer poolMutex.Unlock()
+		if pool.Load() == nil {
+			var err error
+			catalogConnectionString := GetCatalogConnectionStringFromEnv(ctx)
+			config, err := pgxpool.ParseConfig(catalogConnectionString)
+			if err != nil {
+				return shared.CatalogPool{},
+					exceptions.NewCatalogError(fmt.Errorf("unable to parse catalog connection string: %w", err))
+			}
+			config.MaxConns = 3
+			config.MaxConnIdleTime = 90 * time.Second
+			localPool, err := pgxpool.NewWithConfig(ctx, config)
+			if err != nil {
+				return shared.CatalogPool{},
+					exceptions.NewCatalogError(fmt.Errorf("unable to initialize catalog connection pool: %w", err))
+			}
 
-	if err := pool.Ping(ctx); err != nil {
-		return shared.CatalogPool{Pool: pool},
-			exceptions.NewCatalogError(fmt.Errorf("unable to establish connection with catalog: %w", err))
+			if err := localPool.Ping(ctx); err != nil {
+				localPool.Close()
+				return shared.CatalogPool{},
+					exceptions.NewCatalogError(fmt.Errorf("unable to establish connection with catalog: %w", err))
+			}
+			pool.Store(localPool)
+		}
 	}
 
-	return shared.CatalogPool{Pool: pool}, nil
+	return shared.CatalogPool{Pool: pool.Load()}, nil
 }
 
 func GetCatalogConnectionStringFromEnv(ctx context.Context) string {

flow/internal/dynamicconf.go

@@ -862,12 +862,18 @@ func PeerDBS3UuidPrefix(ctx context.Context, env map[string]string) (bool, error
 	return dynamicConfBool(ctx, env, "PEERDB_S3_UUID_PREFIX")
 }
 
+var peerDBS3PartSizeCache CachedDynconfSetting[int64]
+
 func PeerDBS3PartSize(ctx context.Context, env map[string]string) (int64, error) {
-	return dynamicConfSigned[int64](ctx, env, "PEERDB_S3_PART_SIZE")
+	peerDBS3PartSizeCache.InitOnce("PEERDB_S3_PART_SIZE", 30*time.Second, dynamicConfSigned[int64])
+	return peerDBS3PartSizeCache.Get(ctx, env)
 }
 
+var peerDBS3BytesPerAvroFileCache CachedDynconfSetting[int64]
+
 func PeerDBS3BytesPerAvroFile(ctx context.Context, env map[string]string) (int64, error) {
-	return dynamicConfSigned[int64](ctx, env, "PEERDB_S3_BYTES_PER_AVRO_FILE")
+	peerDBS3BytesPerAvroFileCache.InitOnce("PEERDB_S3_BYTES_PER_AVRO_FILE", 30*time.Second, dynamicConfSigned[int64])
+	return peerDBS3BytesPerAvroFileCache.Get(ctx, env)
 }
 
 // Kafka has topic auto create as an option, auto.create.topics.enable

flow/internal/dynamicconf_cache.go

@@ -0,0 +1,60 @@
+package internal
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+)
+
+// CachedDynconfSetting caches the value returned by a typed dynamic setting getter.
+type CachedDynconfSetting[T any] struct {
+	loadedAt time.Time
+	value    T
+	getter   func(context.Context, map[string]string, string) (T, error)
+	name     string
+	ttl      time.Duration
+	once     sync.Once
+	mu       sync.Mutex
+}
+
+// InitOnce configures the setting cache on its first call.
+func (s *CachedDynconfSetting[T]) InitOnce(
+	name string,
+	ttl time.Duration,
+	getter func(context.Context, map[string]string, string) (T, error),
+) {
+	s.once.Do(func() {
+		s.name = name
+		s.ttl = ttl
+		s.getter = getter
+	})
+}
+
+// Get returns the cached value or refreshes it using the configured getter.
+func (s *CachedDynconfSetting[T]) Get(ctx context.Context, env map[string]string) (T, error) {
+	if s.getter == nil {
+		var zero T
+		return zero, errors.New("cached dynamic setting is not initialized")
+	}
+
+	if _, overridden := env[s.name]; overridden {
+		return s.getter(ctx, env, s.name)
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if !s.loadedAt.IsZero() && time.Since(s.loadedAt) < s.ttl {
+		return s.value, nil
+	}
+
+	value, err := s.getter(ctx, nil, s.name)
+	if err != nil {
+		var zero T
+		return zero, err
+	}
+	s.value = value
+	s.loadedAt = time.Now()
+	return value, nil
+}

flow/internal/dynamicconf_cache_test.go

@@ -0,0 +1,165 @@
+package internal
+
+import (
+	"context"
+	"errors"
+	"math/rand/v2"
+	"strconv"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCachedDynconfSettingCachesWithinTTL(t *testing.T) {
+	var calls int
+	getter := func(context.Context, map[string]string, string) (int64, error) {
+		calls++
+		return 10, nil
+	}
+	var setting CachedDynconfSetting[int64]
+	setting.InitOnce("TEST_SETTING", time.Minute, getter)
+
+	for range 2 {
+		value, err := setting.Get(t.Context(), nil)
+		require.NoError(t, err)
+		require.Equal(t, int64(10), value)
+	}
+	require.Equal(t, 1, calls)
+}
+
+func TestCachedDynconfSettingRequiresInitialization(t *testing.T) {
+	var setting CachedDynconfSetting[int64]
+
+	_, err := setting.Get(t.Context(), nil)
+	require.EqualError(t, err, "cached dynamic setting is not initialized")
+}
+
+func TestCachedDynconfSettingRefreshesAfterTTL(t *testing.T) {
+	var calls int64
+	getter := func(context.Context, map[string]string, string) (int64, error) {
+		calls++
+		return calls * 10, nil
+	}
+	var setting CachedDynconfSetting[int64]
+	setting.InitOnce("TEST_SETTING", time.Millisecond, getter)
+
+	value, err := setting.Get(t.Context(), nil)
+	require.NoError(t, err)
+	require.Equal(t, int64(10), value)
+
+	time.Sleep(10 * time.Millisecond)
+
+	value, err = setting.Get(t.Context(), nil)
+	require.NoError(t, err)
+	require.Equal(t, int64(20), value)
+	require.Equal(t, int64(2), calls)
+}
+
+func TestCachedDynconfSettingDoesNotCacheErrors(t *testing.T) {
+	wantErr := errors.New("lookup failed")
+	var calls int
+	getter := func(context.Context, map[string]string, string) (int64, error) {
+		calls++
+		if calls == 1 {
+			return 0, wantErr
+		}
+		return 10, nil
+	}
+	var setting CachedDynconfSetting[int64]
+	setting.InitOnce("TEST_SETTING", time.Minute, getter)
+
+	_, err := setting.Get(t.Context(), nil)
+	require.ErrorIs(t, err, wantErr)
+	value, err := setting.Get(t.Context(), nil)
+	require.NoError(t, err)
+	require.Equal(t, int64(10), value)
+	require.Equal(t, 2, calls)
+}
+
+func TestCachedDynconfSettingCoalescesConcurrentRefreshes(t *testing.T) {
+	const goroutines = 8
+
+	var calls atomic.Int32
+	started := make(chan struct{})
+	release := make(chan struct{})
+	//nolint:unparam // Signature must match the getter type
+	getter := func(context.Context, map[string]string, string) (int64, error) {
+		if calls.Add(1) == 1 {
+			close(started)
+		}
+		<-release
+		return 10, nil
+	}
+	var setting CachedDynconfSetting[int64]
+
+	results := make(chan int64, goroutines)
+	errs := make(chan error, goroutines)
+	var wg sync.WaitGroup
+	wg.Add(goroutines)
+	for range goroutines {
+		go func() {
+			defer wg.Done()
+			setting.InitOnce("TEST_SETTING", time.Minute, getter)
+			value, err := setting.Get(t.Context(), nil)
+			results <- value
+			errs <- err
+		}()
+	}
+
+	<-started
+	close(release)
+	wg.Wait()
+	close(results)
+	close(errs)
+
+	for err := range errs {
+		require.NoError(t, err)
+	}
+	for value := range results {
+		require.Equal(t, int64(10), value)
+	}
+	require.Equal(t, int32(1), calls.Load())
+}
+
+func TestCachedDynconfSettingWithTypedDynconfGetter(t *testing.T) {
+	//nolint:gosec // Test data does not need cryptographically secure randomness.
+	expectedPartSize := rand.Int64()
+	var setting CachedDynconfSetting[int64]
+	setting.InitOnce("PEERDB_S3_PART_SIZE", time.Minute, dynamicConfSigned[int64])
+
+	value, err := setting.Get(t.Context(), map[string]string{
+		"PEERDB_S3_PART_SIZE": strconv.FormatInt(expectedPartSize, 10),
+	})
+	require.NoError(t, err)
+	require.Equal(t, expectedPartSize, value)
+}
+
+func TestCachedDynconfSettingDoesNotCacheEnvOverrides(t *testing.T) {
+	const name = "TEST_SETTING"
+	var calls int
+	getter := func(_ context.Context, env map[string]string, name string) (int64, error) {
+		calls++
+		if value, overridden := env[name]; overridden {
+			return strconv.ParseInt(value, 10, 64)
+		}
+		return 10, nil
+	}
+	var setting CachedDynconfSetting[int64]
+	setting.InitOnce(name, time.Minute, getter)
+
+	value, err := setting.Get(t.Context(), nil)
+	require.NoError(t, err)
+	require.Equal(t, int64(10), value)
+
+	value, err = setting.Get(t.Context(), map[string]string{name: "20"})
+	require.NoError(t, err)
+	require.Equal(t, int64(20), value)
+
+	value, err = setting.Get(t.Context(), nil)
+	require.NoError(t, err)
+	require.Equal(t, int64(10), value)
+	require.Equal(t, 2, calls)
+}