snapshot_activity.go

package activities

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"go.temporal.io/sdk/activity"

	"github.com/PeerDB-io/peerdb/flow/alerting"
	"github.com/PeerDB-io/peerdb/flow/connectors"
	"github.com/PeerDB-io/peerdb/flow/generated/protos"
	"github.com/PeerDB-io/peerdb/flow/internal"
	"github.com/PeerDB-io/peerdb/flow/pkg/common"
	"github.com/PeerDB-io/peerdb/flow/shared"
)

type SlotSnapshotState struct {
	connector      connectors.CDCPullConnectorCore
	connectorClose func(ctx context.Context)
	slotConn       interface{ Close(context.Context) error }
	snapshotName   string
}

type TxSnapshotState struct {
	SnapshotName        string
	SnapshotStagingPath string
}

type SnapshotActivity struct {
	Alerter             *alerting.Alerter
	CatalogPool         shared.CatalogPool
	SlotSnapshotStates  map[string]SlotSnapshotState
	TxSnapshotStates    map[string]TxSnapshotState
	SnapshotStatesMutex sync.Mutex
}

// closes the slot signal
func (a *SnapshotActivity) CloseSlotKeepAlive(ctx context.Context, flowJobName string) error {
	a.SnapshotStatesMutex.Lock()
	defer a.SnapshotStatesMutex.Unlock()

	if s, ok := a.SlotSnapshotStates[flowJobName]; ok {
		if s.slotConn != nil {
			logger := internal.LoggerFromCtx(ctx)
			if err := s.slotConn.Close(ctx); err != nil {
				logger.Error("failed to close the slot connection", slog.Any("error", err))
			} else {
				logger.Info("closed the slot connection")
			}
		}
		s.connectorClose(ctx)
		delete(a.SlotSnapshotStates, flowJobName)
	}
	a.Alerter.EmitFlowInfoTelemetryEvent(ctx, flowJobName, "Ended Snapshot Flow Job")

	return nil
}

func (a *SnapshotActivity) SetupReplication(
	ctx context.Context,
	config *protos.SetupReplicationInput,
) (*protos.SetupReplicationOutput, error) {
	ctx = context.WithValue(ctx, shared.FlowNameKey, config.FlowJobName)
	logger := internal.LoggerFromCtx(ctx)
	a.Alerter.LogFlowInfo(ctx, config.FlowJobName, "Setting up replication slot and publication")
	a.Alerter.EmitFlowInfoTelemetryEvent(ctx, config.FlowJobName, "Started Snapshot Flow Job")

	conn, connClose, err := connectors.GetByNameAs[connectors.CDCPullConnectorCore](ctx, nil, a.CatalogPool, config.PeerName)
	if err != nil {
		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("failed to get connector: %w", err))
	}

	logger.Info("waiting for slot to be created...")
	slotInfo, err := conn.SetupReplication(ctx, config)

	if err != nil {
		connClose(ctx)
		// it is important to close the connection here as it is not closed in CloseSlotKeepAlive
		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, shared.WrapError("slot error", err))
	} else if slotInfo.Conn == nil && slotInfo.SlotName == "" {
		connClose(ctx)
		logger.Info("replication setup without slot")
		return nil, nil
	} else {
		logger.Info("slot created", slog.String("SlotName", slotInfo.SlotName))
	}

	a.SnapshotStatesMutex.Lock()
	a.SlotSnapshotStates[config.FlowJobName] = SlotSnapshotState{
		slotConn:       slotInfo.Conn,
		snapshotName:   slotInfo.SnapshotName,
		connector:      conn,
		connectorClose: connClose,
	}
	a.SnapshotStatesMutex.Unlock()

	a.Alerter.LogFlowInfo(ctx, config.FlowJobName, "Replication slot and publication setup complete")

	return &protos.SetupReplicationOutput{
		SlotName:     slotInfo.SlotName,
		SnapshotName: slotInfo.SnapshotName,
	}, nil
}

func (a *SnapshotActivity) MaintainTx(ctx context.Context, sessionID string, flowName string, peer string, env map[string]string) error {
	shutdown := common.HeartbeatRoutine(ctx, func() string {
		return "maintaining transaction snapshot"
	})
	defer shutdown()
	conn, connClose, err := connectors.GetByNameAs[connectors.CDCPullConnectorCore](ctx, nil, a.CatalogPool, peer)
	if err != nil {
		return a.Alerter.LogFlowError(ctx, sessionID, err)
	}
	defer connClose(ctx)

	exportSnapshotOutput, tx, err := conn.ExportTxSnapshot(ctx, flowName, env)
	if err != nil {
		return err
	}

	a.SnapshotStatesMutex.Lock()
	if exportSnapshotOutput != nil {
		a.TxSnapshotStates[sessionID] = TxSnapshotState{
			SnapshotName:        exportSnapshotOutput.SnapshotName,
			SnapshotStagingPath: exportSnapshotOutput.SnapshotStagingPath,
		}
	} else {
		a.TxSnapshotStates[sessionID] = TxSnapshotState{}
	}
	a.SnapshotStatesMutex.Unlock()

	logger := internal.LoggerFromCtx(ctx)
	start := time.Now()
	for {
		logger.Info("maintaining export snapshot transaction", slog.Int64("seconds", int64(time.Since(start).Round(time.Second)/time.Second)))
		if ctx.Err() != nil {
			a.SnapshotStatesMutex.Lock()
			delete(a.TxSnapshotStates, sessionID)
			a.SnapshotStatesMutex.Unlock()
			if err := conn.FinishExport(tx); err != nil {
				logger.Error("finish export error", slog.Any("error", err))
				return err
			}
			return nil
		}
		time.Sleep(time.Minute)
	}
}

func (a *SnapshotActivity) WaitForExportSnapshot(ctx context.Context, sessionID string) (*TxSnapshotState, error) {
	logger := internal.LoggerFromCtx(ctx)
	attempt := 0
	for {
		a.SnapshotStatesMutex.Lock()
		tsc, ok := a.TxSnapshotStates[sessionID]
		a.SnapshotStatesMutex.Unlock()
		if ok {
			return &tsc, nil
		}
		activity.RecordHeartbeat(ctx, "wait another second for snapshot export")
		attempt += 1
		if attempt > 2 {
			logger.Info("waiting on snapshot export", slog.Int("attempt", attempt))
		}
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		time.Sleep(time.Second)
	}
}

func (a *SnapshotActivity) LoadTableSchema(
	ctx context.Context,
	flowName string,
	tableName string,
) (*protos.TableSchema, error) {
	return internal.LoadTableSchemaFromCatalog(ctx, a.CatalogPool, flowName, tableName)
}

func (a *SnapshotActivity) GetPeerType(ctx context.Context, name string) (protos.DBType, error) {
	return connectors.LoadPeerType(ctx, a.CatalogPool, name)
}

func (a *SnapshotActivity) GetDefaultPartitionKeyForTables(
	ctx context.Context,
	input *protos.FlowConnectionConfigsCore,
) (*protos.GetDefaultPartitionKeyForTablesOutput, error) {
	conn, connClose, err := connectors.GetByNameAs[connectors.QRepPullConnectorCore](ctx, nil, a.CatalogPool, input.SourceName)
	if err != nil {
		return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to get connector: %w", err))
	}
	defer connClose(ctx)

	output, err := conn.GetDefaultPartitionKeyForTables(ctx, &protos.GetDefaultPartitionKeyForTablesInput{
		TableMappings: input.TableMappings,
	})
	if err != nil {
		return nil, a.Alerter.LogFlowError(ctx, input.FlowJobName, fmt.Errorf("failed to check if tables can parallel load: %w", err))
	}

	return output, nil
}