fix

Author: serichoi65

pkg/postgres/migrations/202511171438_withdrawalAndDeallocationQueues/up.go

@@ -15,12 +15,11 @@ func (m *Migration) Up(db *sql.DB, grm *gorm.DB, cfg *config.Config) error {
 	queries := []string{
 
 		// =============================================================================
-		// PART 4: Create withdrawal_queue_share_snapshots table for rewards calculation
+		// PART 1: Create withdrawal_queue_share_snapshots table for rewards calculation
 		// =============================================================================
 
 		// Table to track shares in withdrawal queue that should still earn rewards
 		// Stakers continue earning while in 14-day queue because they're still taking slashing risk
-		// Follows naming convention: *_snapshots suffix for snapshot tables (consistent with staker_share_snapshots, etc.)
 		`create table if not exists withdrawal_queue_share_snapshots (
 			staker varchar not null,
 			strategy varchar not null,
@@ -39,7 +38,7 @@ func (m *Migration) Up(db *sql.DB, grm *gorm.DB, cfg *config.Config) error {
 		`create index if not exists idx_withdrawal_queue_share_snapshots_staker on withdrawal_queue_share_snapshots(staker, snapshot)`,
 
 		// =============================================================================
-		// PART 5: Create deallocation_queue_snapshots table for operator allocation tracking
+		// PART 2: Create deallocation_queue_snapshots table for operator allocation tracking
 		// =============================================================================
 
 		// Table to track operator allocation decreases that haven't reached effective_date yet

pkg/rewards/deallocationQueueShareSnapshots.go

@@ -42,7 +42,7 @@ const deallocationQueueShareSnapshotsQuery = `
 		from operator_allocations oa
 		inner join blocks b on oa.block_number = b.number
 		where
-			oa.effective_date is not null  -- Only Sabine fork records
+			oa.effective_date is not null
 			-- Allocation already recorded before or on snapshot
 			and date(b.block_time) <= '{{.snapshotDate}}'
 			-- Effective date is in the future relative to snapshot
@@ -85,8 +85,6 @@ const deallocationQueueShareSnapshotsQuery = `
 	on conflict on constraint uniq_deallocation_queue_snapshots do nothing;
 `
 
-// DeallocationQueueSnapshot represents operator allocation decreases that haven't
-// reached their effective_date yet, and should still be counted for rewards.
 type DeallocationQueueSnapshot struct {
 	Operator          string `gorm:"column:operator;primaryKey"`
 	Avs               string `gorm:"column:avs;primaryKey"`
@@ -101,13 +99,7 @@ func (DeallocationQueueSnapshot) TableName() string {
 	return "deallocation_queue_snapshots"
 }
 
-// GenerateAndInsertDeallocationQueueSnapshots calculates and inserts operator allocation
-// decreases that should still be counted for rewards because their effective_date hasn't
-// been reached yet.
-//
-// This feature is only active after the Sabine fork (when effective_date is populated).
 func (r *RewardsCalculator) GenerateAndInsertDeallocationQueueSnapshots(snapshotDate string) error {
-	// Check if we're past the Sabine fork
 	forks, err := r.globalConfig.GetRewardsSqlForkDates()
 	if err != nil {
 		r.logger.Sugar().Errorw("Failed to get rewards fork dates", "error", err)
@@ -120,7 +112,6 @@ func (r *RewardsCalculator) GenerateAndInsertDeallocationQueueSnapshots(snapshot
 		return nil
 	}
 
-	// Get the block number for the snapshot date to compare with fork block
 	var maxBlock uint64
 	res := r.grm.Raw(`
 		SELECT COALESCE(MAX(number), 0) as max_block
@@ -133,7 +124,6 @@ func (r *RewardsCalculator) GenerateAndInsertDeallocationQueueSnapshots(snapshot
 		return res.Error
 	}
 
-	// Only apply deallocation queue logic if we're past the Sabine fork
 	if maxBlock < sabineFork.BlockNumber {
 		r.logger.Sugar().Debugw("Snapshot date is before Sabine fork, skipping deallocation queue logic",
 			zap.String("snapshotDate", snapshotDate),
@@ -168,38 +158,41 @@ func (r *RewardsCalculator) GenerateAndInsertDeallocationQueueSnapshots(snapshot
 	return nil
 }
 
-// ListDeallocationQueueSnapshots returns all deallocation queue snapshots for debugging
-func (r *RewardsCalculator) ListDeallocationQueueSnapshots() ([]*DeallocationQueueSnapshot, error) {
-	var snapshots []*DeallocationQueueSnapshot
-	res := r.grm.Model(&DeallocationQueueSnapshot{}).Find(&snapshots)
-	if res.Error != nil {
-		r.logger.Sugar().Errorw("Failed to list deallocation queue snapshots", "error", res.Error)
-		return nil, res.Error
+func (r *RewardsCalculator) AdjustOperatorShareSnapshotsForDeallocationQueue(snapshotDate string) error {
+	adjustQuery := `
+	insert into operator_share_snapshots(operator, strategy, shares, snapshot)
+		select
+			dqs.operator,
+			dqs.strategy,
+			sum(dqs.magnitude_decrease::numeric)::text as shares,
+			dqs.snapshot as snapshot
+		from deallocation_queue_snapshots dqs
+		where dqs.snapshot = '{{.snapshotDate}}'
+		group by dqs.operator, dqs.strategy, dqs.snapshot
+		on conflict on constraint uniq_operator_share_snapshots
+		do update set shares = operator_share_snapshots.shares + EXCLUDED.shares;`
+
+	query, err := rewardsUtils.RenderQueryTemplate(adjustQuery, map[string]interface{}{
+		"snapshotDate": snapshotDate,
+	})
+	if err != nil {
+		r.logger.Sugar().Errorw("Failed to render deallocation queue adjustment query template", "error", err)
+		return err
 	}
-	return snapshots, nil
-}
 
-// adjustOperatorAllocationSnapshotsForDeallocationQueueQuery adds deallocation queue
-// allocations back to operator allocation snapshots.
-//
-// This is similar to withdrawal queue adjustment but for operator allocations.
-// We add back the magnitude_decrease to the current allocation to get the pre-deallocation value.
-const adjustOperatorAllocationSnapshotsForDeallocationQueueQuery = `
-	-- NOT YET IMPLEMENTED
-	-- This would need to integrate with operator allocation snapshots used in rewards v2.2
-	-- For now, this is a placeholder for future operator set rewards calculation
-	-- TODO: Integrate with operator set rewards calculation
-`
+	res := r.grm.Debug().Exec(query)
+	if res.Error != nil {
+		r.logger.Sugar().Errorw("Failed to adjust operator_share_snapshots for deallocation queue",
+			zap.String("snapshotDate", snapshotDate),
+			zap.Error(res.Error),
+		)
+		return res.Error
+	}
 
-// AdjustOperatorAllocationSnapshotsForDeallocationQueue adds deallocation queue
-// allocations back to operator allocation snapshots.
-//
-// NOTE: This is a placeholder for integration with operator set rewards (v2.2).
-// The actual adjustment logic will depend on how operator allocations are used in rewards.
-func (r *RewardsCalculator) AdjustOperatorAllocationSnapshotsForDeallocationQueue(snapshotDate string) error {
-	// TODO: Implement adjustment logic once operator set rewards calculation structure is finalized
-	r.logger.Sugar().Debugw("Deallocation queue adjustment not yet implemented",
+	r.logger.Sugar().Infow("Adjusted operator share snapshots for deallocation queue",
 		zap.String("snapshotDate", snapshotDate),
+		zap.Int64("rowsAffected", res.RowsAffected),
 	)
+
 	return nil
 }

pkg/rewards/rewards.go

@@ -674,20 +674,24 @@ func (rc *RewardsCalculator) generateSnapshotData(snapshotDate string) error {
 	}
 	rc.logger.Sugar().Debugw("Generated operator share snapshots")
 
-	// Generate deallocation queue snapshots - operators continue earning on old allocation until effective_date
 	if err = rc.GenerateAndInsertDeallocationQueueSnapshots(snapshotDate); err != nil {
 		rc.logger.Sugar().Errorw("Failed to generate deallocation queue snapshots", "error", err)
 		return err
 	}
 	rc.logger.Sugar().Debugw("Generated deallocation queue snapshots")
 
+	if err = rc.AdjustOperatorShareSnapshotsForDeallocationQueue(snapshotDate); err != nil {
+		rc.logger.Sugar().Errorw("Failed to adjust operator share snapshots for deallocation queue", "error", err)
+		return err
+	}
+	rc.logger.Sugar().Debugw("Adjusted operator share snapshots for deallocation queue")
+
 	if err = rc.GenerateAndInsertStakerShareSnapshots(snapshotDate); err != nil {
 		rc.logger.Sugar().Errorw("Failed to generate staker share snapshots", "error", err)
 		return err
 	}
 	rc.logger.Sugar().Debugw("Generated staker share snapshots")
 
-	// Adjust staker share snapshots to include withdrawal queue shares
 	if err = rc.AdjustStakerShareSnapshotsForWithdrawalQueue(snapshotDate); err != nil {
 		rc.logger.Sugar().Errorw("Failed to adjust staker share snapshots for withdrawal queue", "error", err)
 		return err

pkg/rewards/withdrawalQueueShareSnapshots.go

@@ -73,10 +73,6 @@ func (WithdrawalQueueShareSnapshot) TableName() string {
 	return "withdrawal_queue_share_snapshots"
 }
 
-// GenerateAndInsertWithdrawalQueueShares calculates and inserts shares in withdrawal queue
-// that should still earn rewards for the given snapshot date.
-//
-// This feature is only active after the Sabine fork.
 func (r *RewardsCalculator) GenerateAndInsertWithdrawalQueueShares(snapshotDate string) error {
 	forks, err := r.globalConfig.GetRewardsSqlForkDates()
 	if err != nil {
@@ -90,7 +86,6 @@ func (r *RewardsCalculator) GenerateAndInsertWithdrawalQueueShares(snapshotDate
 		return nil
 	}
 
-	// Get the block number for the snapshot date to compare with fork block
 	var maxBlock uint64
 	res := r.grm.Raw(`
 		SELECT COALESCE(MAX(number), 0) as max_block
@@ -136,45 +131,20 @@ func (r *RewardsCalculator) GenerateAndInsertWithdrawalQueueShares(snapshotDate
 	return nil
 }
 
-// ListWithdrawalQueueShareSnapshots returns all withdrawal queue share snapshots for debugging
-func (r *RewardsCalculator) ListWithdrawalQueueShareSnapshots() ([]*WithdrawalQueueShareSnapshot, error) {
-	var snapshots []*WithdrawalQueueShareSnapshot
-	res := r.grm.Model(&WithdrawalQueueShareSnapshot{}).Find(&snapshots)
-	if res.Error != nil {
-		r.logger.Sugar().Errorw("Failed to list withdrawal queue share snapshots", "error", res.Error)
-		return nil, res.Error
-	}
-	return snapshots, nil
-}
-
-// adjustStakerShareSnapshotsForWithdrawalQueueQuery adds withdrawal queue shares
-// to existing staker_share_snapshots for the given snapshot date.
-//
-// This is executed AFTER GenerateAndInsertStakerShareSnapshots to add back shares
-// that are in the withdrawal queue and should still earn rewards.
-const adjustStakerShareSnapshotsForWithdrawalQueueQuery = `
-	-- Add withdrawal queue shares to existing staker share snapshots
-	-- If a staker/strategy already exists, add to their shares
-	-- Otherwise, insert a new row
-	insert into staker_share_snapshots(staker, strategy, shares, snapshot)
-	select
-		wqss.staker,
-		wqss.strategy,
-		wqss.shares,
-		wqss.snapshot as snapshot
-	from withdrawal_queue_share_snapshots wqss
-	where wqss.snapshot = '{{.snapshotDate}}'
-	on conflict on constraint uniq_staker_share_snapshots
-	do update set
-		-- Add withdrawal queue shares to existing snapshot shares
-		shares = staker_share_snapshots.shares + EXCLUDED.shares;
-`
-
-// AdjustStakerShareSnapshotsForWithdrawalQueue adds withdrawal queue shares to
-// staker share snapshots. This ensures stakers continue earning rewards while
-// their withdrawals are in the 14-day queue.
 func (r *RewardsCalculator) AdjustStakerShareSnapshotsForWithdrawalQueue(snapshotDate string) error {
-	query, err := rewardsUtils.RenderQueryTemplate(adjustStakerShareSnapshotsForWithdrawalQueueQuery, map[string]interface{}{
+	adjustQuery := `
+	insert into staker_share_snapshots(staker, strategy, shares, snapshot)
+		select
+			wqss.staker,
+			wqss.strategy,
+			wqss.shares,
+			wqss.snapshot as snapshot
+		from withdrawal_queue_share_snapshots wqss
+		where wqss.snapshot = '{{.snapshotDate}}'
+		on conflict on constraint uniq_staker_share_snapshots
+		do update set shares = staker_share_snapshots.shares + EXCLUDED.shares;`
+
+	query, err := rewardsUtils.RenderQueryTemplate(adjustQuery, map[string]interface{}{
 		"snapshotDate": snapshotDate,
 	})
 	if err != nil {