fix: server panics during shutdown with reporting metrics: failed to …

…store jobs: context canceled (#4228)

enterprise/reporting/error_index/error_index_reporting.go

@@ -112,7 +112,7 @@ func NewErrorIndexReporter(ctx context.Context, log logger.Logger, configSubscri
 }
 
 // Report reports the metrics to the errorIndex JobsDB
-func (eir *ErrorIndexReporter) Report(metrics []*types.PUReportedMetric, tx *Tx) error {
+func (eir *ErrorIndexReporter) Report(ctx context.Context, metrics []*types.PUReportedMetric, tx *Tx) error {
  failedAt := eir.now()
 
  var jobs []*jobsdb.JobT
@@ -171,8 +171,8 @@ func (eir *ErrorIndexReporter) Report(metrics []*types.PUReportedMetric, tx *Tx)
  if err != nil {
  return fmt.Errorf("failed to resolve jobsdb: %w", err)
  }
- if err := db.WithStoreSafeTxFromTx(eir.ctx, tx, func(tx jobsdb.StoreSafeTx) error {
- return db.StoreInTx(eir.ctx, tx, jobs)
+ if err := db.WithStoreSafeTxFromTx(ctx, tx, func(tx jobsdb.StoreSafeTx) error {
+ return db.StoreInTx(ctx, tx, jobs)
  }); err != nil {
  return fmt.Errorf("failed to store jobs: %w", err)
  }

enterprise/reporting/error_index/error_index_reporting_test.go

@@ -242,7 +242,7 @@ func TestErrorIndexReporter(t *testing.T) {
  sqlTx, err := postgresContainer.DB.Begin()
  require.NoError(t, err)
  tx := &Tx{Tx: sqlTx}
- err = eir.Report(tc.reports, tx)
+ err = eir.Report(context.Background(), tc.reports, tx)
  require.NoError(t, err)
  require.NoError(t, tx.Commit())
  db, err := eir.resolveJobsDB(tx)
@@ -306,7 +306,7 @@ func TestErrorIndexReporter(t *testing.T) {
  sqlTx, err := postgresContainer.DB.Begin()
  require.NoError(t, err)
  tx := &Tx{Tx: sqlTx}
- err = eir.Report([]*types.PUReportedMetric{}, tx)
+ err = eir.Report(context.Background(), []*types.PUReportedMetric{}, tx)
  require.NoError(t, err)
  require.NoError(t, tx.Commit())
 
@@ -342,7 +342,7 @@ func TestErrorIndexReporter(t *testing.T) {
  sqlTx, err := pg2.DB.Begin()
  require.NoError(t, err)
  tx := &Tx{Tx: sqlTx}
- err = eir.Report([]*types.PUReportedMetric{
+ err = eir.Report(context.Background(), []*types.PUReportedMetric{
  {
  ConnectionDetails: types.ConnectionDetails{
  SourceID: sourceID,
@@ -409,7 +409,7 @@ func TestErrorIndexReporter(t *testing.T) {
  sqlTx, err := pg1.DB.Begin()
  require.NoError(t, err)
  tx := &Tx{Tx: sqlTx}
- err = eir.Report([]*types.PUReportedMetric{
+ err = eir.Report(context.Background(), []*types.PUReportedMetric{
  {
  ConnectionDetails: types.ConnectionDetails{
  SourceID: sourceID,
@@ -443,7 +443,7 @@ func TestErrorIndexReporter(t *testing.T) {
  sqlTx, err := pg3.DB.Begin()
  require.NoError(t, err)
  tx := &Tx{Tx: sqlTx}
- err = eir.Report([]*types.PUReportedMetric{
+ err = eir.Report(context.Background(), []*types.PUReportedMetric{
  {
  ConnectionDetails: types.ConnectionDetails{
  SourceID: sourceID,
@@ -571,7 +571,7 @@ func TestErrorIndexReporter(t *testing.T) {
  require.NoError(t, err)
 
  tx := &Tx{Tx: sqlTx}
- err = eir.Report(reports, tx)
+ err = eir.Report(context.Background(), reports, tx)
  require.NoError(t, err)
  require.NoError(t, tx.Commit())
 

enterprise/reporting/error_reporting.go

@@ -158,13 +158,13 @@ func (edr *ErrorDetailReporter) GetSyncer(syncerKey string) *types.SyncSource {
  return edr.syncers[syncerKey]
 }
 
-func (edr *ErrorDetailReporter) Report(metrics []*types.PUReportedMetric, txn *Tx) error {
+func (edr *ErrorDetailReporter) Report(ctx context.Context, metrics []*types.PUReportedMetric, txn *Tx) error {
  edr.log.Debug("[ErrorDetailReport] Report method called\n")
  if len(metrics) == 0 {
  return nil
  }
 
- stmt, err := txn.Prepare(pq.CopyIn(ErrorDetailReportsTable, ErrorDetailReportsColumns...))
+ stmt, err := txn.PrepareContext(ctx, pq.CopyIn(ErrorDetailReportsTable, ErrorDetailReportsColumns...))
  if err != nil {
  edr.log.Errorf("Failed during statement preparation: %v", err)
  return fmt.Errorf("preparing statement: %v", err)
@@ -180,7 +180,7 @@ func (edr *ErrorDetailReporter) Report(metrics []*types.PUReportedMetric, txn *T
 
  // extract error-message & error-code
  errDets := edr.extractErrorDetails(metric.StatusDetail.SampleResponse)
- _, err = stmt.Exec(
+ _, err = stmt.ExecContext(ctx,
  workspaceID,
  edr.namespace,
  edr.instanceID,
@@ -203,7 +203,7 @@ func (edr *ErrorDetailReporter) Report(metrics []*types.PUReportedMetric, txn *T
  }
  }
 
- _, err = stmt.Exec()
+ _, err = stmt.ExecContext(ctx)
  if err != nil {
  edr.log.Errorf("Failed during statement preparation: %v", err)
  return fmt.Errorf("executing final statement: %v", err)

enterprise/reporting/event_stats.go

@@ -1,6 +1,7 @@
 package reporting
 
 import (
+ "context"
  "strconv"
 
  "github.com/rudderlabs/rudder-go-kit/stats"
@@ -39,7 +40,7 @@ func (es *EventStatsReporter) Record(metrics []*types.PUReportedMetric) {
  }
 }
 
-func (es *EventStatsReporter) Report(metrics []*types.PUReportedMetric, tx *Tx) error {
+func (es *EventStatsReporter) Report(_ context.Context, metrics []*types.PUReportedMetric, tx *Tx) error {
  tx.AddSuccessListener(func() {
  es.Record(metrics)
  })

enterprise/reporting/mediator.go

@@ -70,9 +70,9 @@ func NewReportingMediator(ctx context.Context, log logger.Logger, enterpriseToke
  return rm
 }
 
-func (rm *Mediator) Report(metrics []*types.PUReportedMetric, txn *Tx) error {
+func (rm *Mediator) Report(ctx context.Context, metrics []*types.PUReportedMetric, txn *Tx) error {
  for _, reporter := range rm.reporters {
- if err := reporter.Report(metrics, txn); err != nil {
+ if err := reporter.Report(ctx, metrics, txn); err != nil {
  return err
  }
  }

enterprise/reporting/noop.go

@@ -1,14 +1,16 @@
 package reporting
 
 import (
+ "context"
+
  . "github.com/rudderlabs/rudder-server/utils/tx" //nolint:staticcheck
  "github.com/rudderlabs/rudder-server/utils/types"
 )
 
 // NOOP reporting implementation that does nothing
 type NOOP struct{}
 
-func (*NOOP) Report(_ []*types.PUReportedMetric, _ *Tx) error {
+func (*NOOP) Report(_ context.Context, _ []*types.PUReportedMetric, _ *Tx) error {
  return nil
 }
 

enterprise/reporting/reporting.go

@@ -542,12 +542,12 @@ func transformMetricForPII(metric types.PUReportedMetric, piiColumns []string) t
  return metric
 }
 
-func (r *DefaultReporter) Report(metrics []*types.PUReportedMetric, txn *Tx) error {
+func (r *DefaultReporter) Report(ctx context.Context, metrics []*types.PUReportedMetric, txn *Tx) error {
  if len(metrics) == 0 {
  return nil
  }
 
- stmt, err := txn.Prepare(pq.CopyIn(ReportsTable,
+ stmt, err := txn.PrepareContext(ctx, pq.CopyIn(ReportsTable,
  "workspace_id", "namespace", "instance_id",
  "source_definition_id",
  "source_category",
@@ -617,7 +617,7 @@ func (r *DefaultReporter) Report(metrics []*types.PUReportedMetric, txn *Tx) err
  return fmt.Errorf("executing statement: %v", err)
  }
  }
- if _, err = stmt.Exec(); err != nil {
+ if _, err = stmt.ExecContext(ctx); err != nil {
  return fmt.Errorf("executing final statement: %v", err)
  }
 

processor/processor.go

@@ -2222,13 +2222,13 @@ func (proc *Handle) Store(partition string, in *storeMessage) {
  if err != nil {
  return fmt.Errorf("publishing rsources stats: %w", err)
  }
- err = proc.saveDroppedJobs(in.droppedJobs, tx.Tx())
+ err = proc.saveDroppedJobs(ctx, in.droppedJobs, tx.Tx())
  if err != nil {
  return fmt.Errorf("saving dropped jobs: %w", err)
  }
 
  if proc.isReportingEnabled() {
- if err = proc.reporting.Report(in.reportMetrics, tx.Tx()); err != nil {
+ if err = proc.reporting.Report(ctx, in.reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }
@@ -2684,7 +2684,7 @@ func (proc *Handle) transformSrcDest(
  }
 }
 
-func (proc *Handle) saveDroppedJobs(droppedJobs []*jobsdb.JobT, tx *Tx) error {
+func (proc *Handle) saveDroppedJobs(ctx context.Context, droppedJobs []*jobsdb.JobT, tx *Tx) error {
  if len(droppedJobs) > 0 {
  for i := range droppedJobs { // each dropped job should have a unique jobID in the scope of the batch
  droppedJobs[i].JobID = int64(i)
@@ -2694,7 +2694,7 @@ func (proc *Handle) saveDroppedJobs(droppedJobs []*jobsdb.JobT, tx *Tx) error {
  rsources.IgnoreDestinationID(),
  )
  rsourcesStats.JobsDropped(droppedJobs)
- return rsourcesStats.Publish(context.TODO(), tx.Tx)
+ return rsourcesStats.Publish(ctx, tx.Tx)
  }
  return nil
 }

router/batchrouter/handle.go

@@ -763,7 +763,7 @@ func (brt *Handle) updateJobStatus(batchJobs *BatchedJobs, isWarehouse bool, err
  }
 
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
+ if err = brt.reporting.Report(ctx, reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }

router/batchrouter/handle_async.go

@@ -58,7 +58,7 @@ func (brt *Handle) updateJobStatuses(ctx context.Context, destinationID string,
  }
 
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
+ if err = brt.reporting.Report(ctx, reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }
@@ -688,7 +688,7 @@ func (brt *Handle) setMultipleJobStatus(asyncOutput common.AsyncUploadOutput, at
  }
 
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
+ if err = brt.reporting.Report(ctx, reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }

router/batchrouter/worker.go

@@ -142,7 +142,7 @@ func (w *worker) processJobAsync(jobsWg *sync.WaitGroup, destinationJobs *Destin
  return fmt.Errorf("marking %s job statuses as aborted: %w", brt.destType, err)
  }
  if brt.reporting != nil && brt.reportingEnabled {
- if err = brt.reporting.Report(reportMetrics, tx.Tx()); err != nil {
+ if err = brt.reporting.Report(ctx, reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  }

router/factory.go

@@ -1,6 +1,8 @@
 package router
 
 import (
+ "context"
+
  "github.com/rudderlabs/rudder-go-kit/config"
  "github.com/rudderlabs/rudder-go-kit/logger"
  backendconfig "github.com/rudderlabs/rudder-server/backend-config"
@@ -50,5 +52,5 @@ func (f *Factory) New(destination *backendconfig.DestinationT) *Handle {
 }
 
 type reporter interface {
- Report(metrics []*utilTypes.PUReportedMetric, txn *Tx) error
+ Report(ctx context.Context, metrics []*utilTypes.PUReportedMetric, txn *Tx) error
 }

router/handle.go

@@ -423,7 +423,7 @@ func (rt *Handle) commitStatusList(workerJobStatuses *[]workerJobStatus) {
  if err != nil {
  return err
  }
- if err = rt.Reporting.Report(reportMetrics, tx.Tx()); err != nil {
+ if err = rt.Reporting.Report(ctx, reportMetrics, tx.Tx()); err != nil {
  return fmt.Errorf("reporting metrics: %w", err)
  }
  return nil

utils/types/types.go

@@ -3,6 +3,7 @@
 package types
 
 import (
+ "context"
  "net/http"
  "time"
 
@@ -62,7 +63,7 @@ type ConfigEnvI interface {
 // Reporting is interface to report metrics
 type Reporting interface {
  // Report reports metrics to reporting service
- Report(metrics []*PUReportedMetric, tx *Tx) error
+ Report(ctx context.Context, metrics []*PUReportedMetric, tx *Tx) error
 
  // DatabaseSyncer creates reporting tables in the database and returns a function to periodically sync the data
  DatabaseSyncer(c SyncerConfig) ReportingSyncer

warehouse/router/upload.go

@@ -522,6 +522,7 @@ func (job *UploadJob) setUploadStatus(statusOpts UploadStatusOpts) (err error) {
  }
  if job.config.reportingEnabled {
  err = job.reporting.Report(
+ job.ctx,
  []*types.PUReportedMetric{&statusOpts.ReportingMetric},
  tx.Tx,
  )
@@ -717,7 +718,7 @@ func (job *UploadJob) setUploadError(statusError error, state string) (string, e
  })
  }
  if job.config.reportingEnabled {
- if err = job.reporting.Report(reportingMetrics, txn.Tx); err != nil {
+ if err = job.reporting.Report(job.ctx, reportingMetrics, txn.Tx); err != nil {
  return "", fmt.Errorf("reporting metrics: %w", err)
  }
  }