improve shutdown handling

Author: 0xSpuddy

cmd/main.go

@@ -132,7 +132,7 @@ func init() {
 	rootCmd.Flags().BoolVar(&testMode, "test", false, "Test mode: verify price feed configurations and calculate medians without starting daemon")
 	rootCmd.Flags().StringVar(&testQueryID, "test-query-id", "", "With --test, only run this custom query id (64-char hex); skips exchange/market tests. Exits non-zero if the query fails.")
 	// Automatic Unbonding flags
-	rootCmd.Flags().Uint32("auto-unbonding-frequency", 0, "Enable automatic unbonding every N days (0 = disabled, 1 - 21 days = valid")
+	rootCmd.Flags().Uint32("auto-unbonding-frequency", 0, "Enable automatic unbonding every N days (0 = disabled, 1 - 21 days = valid)")
 	rootCmd.Flags().Uint32("auto-unbonding-amount", 0, "Amount of tokens in loya to unbond each unbonding transaction (0 = disabled)")
 	rootCmd.Flags().String("auto-unbonding-max-stake-percentage", "0.0", "Maximum percentage of stake to unbond each unbonding transaction (0 = disabled, 1.0 = 100%). If unbonding amount exceeds this percentage, we will skip the unbonding transaction until it exceeds this percentage again.")
 	rootCmd.Flags().Duration("refresh-gas-estimates-interval", 12*time.Hour, "Interval for resetting cached gas estimates and gas-adjustment levels (<=0 disables)")

env.example

@@ -17,6 +17,9 @@ ETH_RPC_URL_FALLBACK=https://eth-mainnet.g.alchemy.com/v2/YOUR_ALCHEMY_API_KEY
 REPORTERS_VALIDATOR_ADDRESS=tellorvaloper1...
 WITHDRAW_FREQUENCY=43200
 
+# Auto-balance is configured with CLI flags:
+# --auto-balance-to-keep, --auto-balance-execution-time, --auto-balance-bridge-to-eth-addr
+
 # Custom query API keys
 # These are used by generated custom_query_config.toml entries that reference ${VAR}.
 CMC_PRO_API_KEY=your_coinmarketcap_api_key

reporter/client/reporter_monitors.go

@@ -229,6 +229,7 @@ func (c *Client) MonitorForTippedQueries(ctx context.Context, wg *sync.WaitGroup
 }
 
 func (c *Client) WithdrawAndStakeEarnedRewardsPeriodically(ctx context.Context, wg *sync.WaitGroup) {
+	defer wg.Done()
 	freqVar := os.Getenv("WITHDRAW_FREQUENCY")
 	if freqVar == "" {
 		freqVar = "43200" // default to being 12 hours or 43200 seconds
@@ -240,20 +241,34 @@ func (c *Client) WithdrawAndStakeEarnedRewardsPeriodically(ctx context.Context,
 	}
 
 	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
 		valAddr := os.Getenv("REPORTERS_VALIDATOR_ADDRESS")
 		if valAddr == "" {
 			fmt.Println("Returning from Withdraw Monitor due to no validator address env variable was found")
-			time.Sleep(time.Duration(frequency) * time.Second)
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(time.Duration(frequency) * time.Second):
+			}
 			continue
 		}
 
 		withdrawMsg := &reportertypes.MsgWithdrawTip{
 			SelectorAddress:  c.accAddr.String(),
 			ValidatorAddress: valAddr,
 		}
-		c.txChan <- TxChannelInfo{Msg: withdrawMsg, isBridge: false, NumRetries: 0, QueryMetaId: 0}
+		c.trySend(ctx, TxChannelInfo{Msg: withdrawMsg, isBridge: false, NumRetries: 0, QueryMetaId: 0})
 
-		time.Sleep(time.Duration(frequency) * time.Second)
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(time.Duration(frequency) * time.Second):
+		}
 	}
 }
 
@@ -279,7 +294,7 @@ func (c *Client) AutoUnbondStakePeriodically(ctx context.Context, wg *sync.WaitG
 	unbondAmount := math.NewInt(int64(amount))
 	valAddr := os.Getenv("REPORTERS_VALIDATOR_ADDRESS")
 	if valAddr == "" {
-		fmt.Println("Returning from Withdraw Monitor due to no validator address env variable was found")
+		fmt.Println("Returning from Auto Unbond Monitor due to no validator address env variable was found")
 		return
 	}
 	for {
@@ -317,7 +332,7 @@ func (c *Client) AutoUnbondStakePeriodically(ctx context.Context, wg *sync.WaitG
 				ValidatorAddress: valAddr,
 				Amount:           sdk.NewCoin("loya", unbondAmount),
 			}
-			c.txChan <- TxChannelInfo{Msg: unbondMsg, isBridge: false, NumRetries: 0, QueryMetaId: 0}
+			c.trySend(ctx, TxChannelInfo{Msg: unbondMsg, isBridge: false, NumRetries: 0, QueryMetaId: 0})
 
 		}
 	}

reporter/client/shutdown_test.go

@@ -78,6 +78,38 @@ func TestStartReporterDaemonTaskLoop_ExitsOnCancelledContext(t *testing.T) {
 	}
 }
 
+func TestWithdrawAndStakeEarnedRewardsPeriodically_ExitsOnCancelledContext(t *testing.T) {
+	c := NewClient(log.NewNopLogger(), "0.001loya")
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		c.WithdrawAndStakeEarnedRewardsPeriodically(ctx, &wg)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("WithdrawAndStakeEarnedRewardsPeriodically did not exit with canceled context")
+	}
+
+	waitDone := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(waitDone)
+	}()
+
+	select {
+	case <-waitDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("WithdrawAndStakeEarnedRewardsPeriodically did not call wg.Done")
+	}
+}
+
 // TestConcurrentTrySendDuringShutdown simulates the real shutdown race:
 // multiple monitor goroutines try to send to txChan while the context is
 // being canceled. None should panic.