feat: implement backoff for snowpipe streaming authorization errors

router/batchrouter/asyncdestinationmanager/snowpipestreaming/channel.go

@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/rudderlabs/rudder-go-kit/logger"
@@ -12,6 +13,11 @@
 	whutils "github.com/rudderlabs/rudder-server/warehouse/utils"
 )
 
+var (
+	errAuthz   = errors.New("snowpipe authorization error")
+	errBackoff = errors.New("snowpipe backoff error")
+)
+
 // initializeChannelWithSchema creates a new channel for the given table if it doesn't exist.
 // If the channel already exists, it checks for new columns and adds them to the table.
 // It returns the channel response after creating or recreating the channel.
@@ -66,7 +72,7 @@
 		snowflakeManager.Cleanup(ctx)
 	}()
 	if err = snowflakeManager.AddColumns(ctx, tableName, columns); err != nil {
-		return fmt.Errorf("adding column: %w", err)
+		return fmt.Errorf("adding column: %w, %w", errAuthz, err)
 	}
 	return nil
 }
@@ -157,10 +163,10 @@
 		snowflakeManager.Cleanup(ctx)
 	}()
 	if err := snowflakeManager.CreateSchema(ctx); err != nil {
-		return nil, fmt.Errorf("creating schema: %w", err)
+		return nil, fmt.Errorf("creating schema: %w, %w", errAuthz, err)
 	}
 	if err := snowflakeManager.CreateTable(ctx, channelReq.TableConfig.Table, eventSchema); err != nil {
-		return nil, fmt.Errorf("creating table: %w", err)
+		return nil, fmt.Errorf("creating table: %w, %w", errAuthz, err)
 	}
 	return m.api.CreateChannel(ctx, channelReq)
 }
@@ -185,7 +191,7 @@
 		snowflakeManager.Cleanup(ctx)
 	}()
 	if err := snowflakeManager.CreateTable(ctx, channelReq.TableConfig.Table, eventSchema); err != nil {
-		return nil, fmt.Errorf("creating table: %w", err)
+		return nil, fmt.Errorf("creating table: %w, %w", errAuthz, err)
 	}
 	return m.api.CreateChannel(ctx, channelReq)
 }
@@ -225,6 +231,9 @@
 }
 
 func (m *Manager) createSnowflakeManager(ctx context.Context, namespace string) (manager.Manager, error) {
+	if m.isInBackoff() {
+		return nil, fmt.Errorf("skipping snowflake manager creation due to backoff: %w", errBackoff)
+	}
 	modelWarehouse := whutils.ModelWarehouse{
 		WorkspaceID: m.destination.WorkspaceID,
 		Destination: *m.destination,
@@ -234,13 +243,5 @@
 	}
 	modelWarehouse.Destination.Config["useKeyPairAuth"] = true // Since we are currently only supporting key pair auth
 
-	sf, err := manager.New(whutils.SnowpipeStreaming, m.appConfig, m.logger, m.statsFactory)
-	if err != nil {
-		return nil, fmt.Errorf("creating snowflake manager: %w", err)
-	}
-	err = sf.Setup(ctx, modelWarehouse, whutils.NewNoOpUploader())
-	if err != nil {
-		return nil, fmt.Errorf("setting up snowflake manager: %w", err)
-	}
-	return sf, nil
+	return m.managerCreator(ctx, modelWarehouse, m.appConfig, m.logger, m.statsFactory)
 }

router/batchrouter/asyncdestinationmanager/snowpipestreaming/snowpipestreaming.go

@@ -4,6 +4,7 @@
 	"bufio"
 	"context"
 	stdjson "encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"os"
@@ -12,6 +13,7 @@
 	"sync"
 	"time"
 
+	"github.com/cenkalti/backoff/v4"
 	"github.com/hashicorp/go-retryablehttp"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/samber/lo"
@@ -31,20 +33,21 @@
 	"github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager/snowpipestreaming/internal/model"
 	"github.com/rudderlabs/rudder-server/utils/misc"
 	"github.com/rudderlabs/rudder-server/utils/timeutil"
+	"github.com/rudderlabs/rudder-server/warehouse/integrations/manager"
 	whutils "github.com/rudderlabs/rudder-server/warehouse/utils"
 )
 
 var json = jsoniter.ConfigCompatibleWithStandardLibrary
 
 func New(
 	conf *config.Config,
-	logger logger.Logger,
+	log logger.Logger,
 	statsFactory stats.Stats,
 	destination *backendconfig.DestinationT,
 ) *Manager {
 	m := &Manager{
 		appConfig: conf,
-		logger: logger.Child("snowpipestreaming").Withn(
+		logger: log.Child("snowpipestreaming").Withn(
 			obskit.WorkspaceID(destination.WorkspaceID),
 			obskit.DestinationID(destination.ID),
 			obskit.DestinationType(destination.DestinationDefinition.Name),
@@ -67,6 +70,9 @@
 	m.config.client.retryMax = conf.GetInt("SnowpipeStreaming.Client.retryMax", 5)
 	m.config.instanceID = conf.GetString("INSTANCE_ID", "1")
 	m.config.maxBufferCapacity = conf.GetReloadableInt64Var(512*bytesize.KB, bytesize.B, "SnowpipeStreaming.maxBufferCapacity")
+	m.config.backoff.initialInterval = conf.GetReloadableDurationVar(1, time.Second, "SnowpipeStreaming.backoffInitialIntervalInSeconds")
+	m.config.backoff.multiplier = conf.GetReloadableFloat64Var(2.0, "SnowpipeStreaming.backoffMultiplier")
+	m.config.backoff.maxInterval = conf.GetReloadableDurationVar(1, time.Hour, "SnowpipeStreaming.backoffMaxIntervalInHours")
 
 	tags := stats.Tags{
 		"module":        "batch_router",
@@ -100,6 +106,17 @@
 		snowpipeapi.New(m.appConfig, m.statsFactory, m.config.client.url, m.requestDoer),
 		destination,
 	)
+	m.managerCreator = func(ctx context.Context, modelWarehouse whutils.ModelWarehouse, conf *config.Config, logger logger.Logger, statsFactory stats.Stats) (manager.Manager, error) {
+		sf, err := manager.New(whutils.SnowpipeStreaming, conf, logger, statsFactory)
+		if err != nil {
+			return nil, fmt.Errorf("creating snowflake manager: %w", err)
+		}
+		err = sf.Setup(ctx, modelWarehouse, whutils.NewNoOpUploader())
+		if err != nil {
+			return nil, fmt.Errorf("setting up snowflake manager: %w", err)
+		}
+		return sf, nil
+	}
 	return m
 }
 
@@ -121,6 +138,10 @@
 	return client
 }
 
+func (m *Manager) Now() time.Time {
+	return m.now()
+}
+
 func (m *Manager) Transform(job *jobsdb.JobT) (string, error) {
 	return common.GetMarshalledData(string(job.EventPayload), job.JobID)
 }
@@ -152,6 +173,12 @@
 
 	discardsChannel, err := m.initializeChannelWithSchema(ctx, asyncDest.Destination.ID, &destConf, discardsTable(), discardsSchema())
 	if err != nil {
+		if errors.Is(err, errAuthz) || errors.Is(err, errBackoff) {
+			if errors.Is(err, errAuthz) {
+				m.setBackOff()
+			}
+			return m.failedJobs(asyncDest, err.Error())
+		}
 		return m.abortJobs(asyncDest, fmt.Errorf("failed to prepare discards channel: %w", err).Error())
 	}
 	m.logger.Infon("Prepared discards channel")
@@ -184,9 +211,17 @@
 		importInfos                   []*importInfo
 		discardImportInfo             *importInfo
 	)
+	shouldResetBackoff := true // backoff should be reset if authz error is not encountered for any of the tables
+	isBackoffSet := false      // should not be set again if already set
 	for _, info := range uploadInfos {
 		imInfo, discardImInfo, err := m.sendEventsToSnowpipe(ctx, asyncDest.Destination.ID, &destConf, info)
 		if err != nil {
+			if errors.Is(err, errAuthz) || errors.Is(err, errBackoff) {
+				shouldResetBackoff = false
+				if errors.Is(err, errAuthz) && !isBackoffSet {
+					m.setBackOff()
+				}
+			}
 			m.logger.Warnn("Failed to send events to Snowpipe",
 				logger.NewStringField("table", info.tableName),
 				obskit.Error(err),
@@ -206,6 +241,9 @@
 			discardImportInfo.Offset = discardImInfo.Offset
 		}
 	}
+	if shouldResetBackoff {
+		m.resetBackoff()
+	}
 	if discardImportInfo != nil {
 		importInfos = append(importInfos, discardImportInfo)
 	}
@@ -245,7 +283,7 @@
 
 	events := make([]*event, 0, eventsCount)
 
-	formattedTS := m.now().Format(misc.RFC3339Milli)
+	formattedTS := m.Now().Format(misc.RFC3339Milli)
 	scanner := bufio.NewScanner(file)
 	scanner.Buffer(nil, int(m.config.maxBufferCapacity.Load()))
 
@@ -289,7 +327,7 @@
 	}
 	log.Infon("Prepared channel", logger.NewStringField("channelID", channelResponse.ChannelID))
 
-	formattedTS := m.now().Format(misc.RFC3339Milli)
+	formattedTS := m.Now().Format(misc.RFC3339Milli)
 	var discardInfos []discardInfo
 	for _, tableEvent := range info.events {
 		discardInfos = append(discardInfos, getDiscardedRecordsFromEvent(tableEvent, channelResponse.SnowpipeSchema, info.tableName, formattedTS)...)
@@ -362,6 +400,16 @@
 	}
 }
 
+func (m *Manager) failedJobs(asyncDest *common.AsyncDestinationStruct, failedReason string) common.AsyncUploadOutput {
+	m.stats.jobs.failed.Count(len(asyncDest.ImportingJobIDs))
+	return common.AsyncUploadOutput{
+		FailedJobIDs:  asyncDest.ImportingJobIDs,
+		FailedCount:   len(asyncDest.ImportingJobIDs),
+		FailedReason:  failedReason,
+		DestinationID: asyncDest.Destination.ID,
+	}
+}
+
 // Poll checks the status of multiple imports using the import ID from pollInput.
 // For the once which have reached the terminal state (success or failure), it caches the import infos in polledImportInfoMap. Later if Poll is called again, it does not need to do the status check again.
 // Once all the imports have reached the terminal state, if any imports have failed, it deletes the channels for those imports.
@@ -549,3 +597,34 @@
 		},
 	}
 }
+
+func (m *Manager) isInBackoff() bool {
+	if m.backoff.next.IsZero() {
+		return false
+	}
+	return m.Now().Before(m.backoff.next)
+}
+
+func (m *Manager) resetBackoff() {
+	m.backoff.next = time.Time{}
+	m.backoff.attempts = 0
+}
+
+func (m *Manager) setBackOff() {
+	b := backoff.NewExponentialBackOff(
+		backoff.WithInitialInterval(m.config.backoff.initialInterval.Load()),
+		backoff.WithMultiplier(m.config.backoff.multiplier.Load()),
+		backoff.WithClockProvider(m),
+		backoff.WithRandomizationFactor(0),
+		backoff.WithMaxElapsedTime(0),
+		backoff.WithMaxInterval(m.config.backoff.maxInterval.Load()),
+	)
+	b.Reset()
+	m.backoff.attempts++
+
+	var d time.Duration
+	for index := int64(0); index < int64(m.backoff.attempts); index++ {
+		d = b.NextBackOff()
+	}
+	m.backoff.next = m.Now().Add(d)
+}