diff --git a/flow/activities/flowable_core.go b/flow/activities/flowable_core.go index 8cda09d87..e00b4a62c 100644 --- a/flow/activities/flowable_core.go +++ b/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 } diff --git a/flow/connectors/utils/monitoring/monitoring.go b/flow/connectors/utils/monitoring/monitoring.go index ddffb4c5c..69baa36ba 100644 --- a/flow/connectors/utils/monitoring/monitoring.go +++ b/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 } diff --git a/flow/internal/catalog.go b/flow/internal/catalog.go index 8c3aa4544..d59a6f279 100644 --- a/flow/internal/catalog.go +++ b/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 { diff --git a/flow/internal/dynamicconf.go b/flow/internal/dynamicconf.go index 57d0d7e26..82e855fd5 100644 --- a/flow/internal/dynamicconf.go +++ b/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 diff --git a/flow/internal/dynamicconf_cache.go b/flow/internal/dynamicconf_cache.go new file mode 100644 index 000000000..5c3b07968 --- /dev/null +++ b/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 +} diff --git a/flow/internal/dynamicconf_cache_test.go b/flow/internal/dynamicconf_cache_test.go new file mode 100644 index 000000000..e83d86c2b --- /dev/null +++ b/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) +}