chore: capture parameter labels in jobsdb update job status stats (#6495)

🔒 Scanned for secrets using gitleaks 8.28.0

# Description

Until now, in `JobsDB#UpdateJobStatus`, all emitted metrics used the
provided `parameterFilters` argument as their labels. However, when
updating job statuses for a mixed set of parameter filters, we had to
pass an empty argument, resulting in metrics without labels.

Now, we are capturing parameter filter stat labels (`source_id` &
`destination_id`) within `JobsDB#UpdateJobStatus` dynamically whenever
the `parameterFilters` argument is empty. To achieve that, the internal
method `doUpdateJobStatusInTx` returns additional statistic counters for
both the number of jobs and the total payload bytes. Thus, the return
type now serves two key purposes:

1. invalidating the `noResultsCache`.
2. Recording `jobsdb_updated_jobs` and `jobsdb_updated_bytes` metrics
with proper labels.

## Linear Ticket

resolves PIPE-2524

## Security

- [x] The code changed/added as part of this pull request won't create
any security issues with how the software is being used.

Author: atzoum

jobsdb/jobsdb.go

@@ -603,7 +603,7 @@ type ParameterFilterT struct {
 }
 
 func (p ParameterFilterT) String() string {
-	return "Name=" + p.Name + ",Value=" + p.Value
+	return p.Name + ":" + p.Value
 }
 
 func (p ParameterFilterT) GetName() string {
@@ -2341,7 +2341,74 @@ func (jd *Handle) getJobsDS(ctx context.Context, ds dataSetT, lastDS bool, param
 	}, true, nil
 }
 
-func (jd *Handle) updateJobStatusDSInTx(ctx context.Context, tx *Tx, ds dataSetT, statusList []*JobStatusT, tags statTags) (updatedStates map[string]map[string]map[ParameterFilterT]struct{}, err error) {
+// updateJobStatusStats is a map containing statistics of job status updates grouped by: workspace -> state -> set of params (stringified) -> stats
+type updateJobStatusStats map[workspaceIDKey]map[jobStateKey]map[parameterFiltersKey]*UpdateJobStatusStats
+
+// workspaceIDKey represents workspace id as key
+type workspaceIDKey string
+
+// jobStateKey represents job state as key (failed, succeeded, etc)
+type jobStateKey string
+
+// parameterFiltersKey represents a list of job parameter filters (stringified) as key
+type parameterFiltersKey string
+
+// Merges metrics from two updateJobStatusStats together
+func (ujss updateJobStatusStats) Merge(other updateJobStatusStats) {
+	for ws, states := range other {
+		if _, ok := ujss[ws]; !ok {
+			ujss[ws] = make(map[jobStateKey]map[parameterFiltersKey]*UpdateJobStatusStats)
+		}
+		for state, paramsMetrics := range states {
+			if _, ok := ujss[ws][state]; !ok {
+				ujss[ws][state] = make(map[parameterFiltersKey]*UpdateJobStatusStats)
+			}
+			for params, metrics := range paramsMetrics {
+				existingMetrics, ok := ujss[ws][state][params]
+				if !ok {
+					existingMetrics = &UpdateJobStatusStats{parameters: metrics.parameters}
+					ujss[ws][state][params] = existingMetrics
+				}
+				existingMetrics.count += metrics.count
+				existingMetrics.bytes += metrics.bytes
+			}
+		}
+	}
+}
+
+// Aggregates metrics by state across all workspaces
+func (ujss updateJobStatusStats) StatsByState() map[jobStateKey]map[parameterFiltersKey]*UpdateJobStatusStats {
+	result := make(map[jobStateKey]map[parameterFiltersKey]*UpdateJobStatusStats)
+	for _, states := range ujss {
+		for state, paramsMetrics := range states {
+			if _, ok := result[state]; !ok {
+				result[state] = make(map[parameterFiltersKey]*UpdateJobStatusStats)
+			}
+			for params, metrics := range paramsMetrics {
+				existingMetrics, ok := result[state][params]
+				if !ok {
+					existingMetrics = &UpdateJobStatusStats{parameters: metrics.parameters}
+					result[state][params] = existingMetrics
+				}
+				existingMetrics.count += metrics.count
+				existingMetrics.bytes += metrics.bytes
+			}
+		}
+	}
+	return result
+}
+
+// Stats for jobs grouped by status and parameters
+type UpdateJobStatusStats struct {
+	// job parameters
+	parameters ParameterFilterList
+	// number of jobs
+	count int
+	// total size of error responses in bytes
+	bytes int
+}
+
+func (jd *Handle) updateJobStatusDSInTx(ctx context.Context, tx *Tx, ds dataSetT, statusList []*JobStatusT, tags statTags) (updatedStates updateJobStatusStats, err error) {
 	if len(statusList) == 0 {
 		return updatedStates, err
 	}
@@ -2350,9 +2417,9 @@ func (jd *Handle) updateJobStatusDSInTx(ctx context.Context, tx *Tx, ds dataSetT
 		"update_job_status_ds_time",
 		&tags,
 	).RecordDuration()()
-	// workspace -> state -> params
-	updatedStates = map[string]map[string]map[ParameterFilterT]struct{}{}
+	updatedStates = updateJobStatusStats{}
 	store := func() error {
+		updatedStates = updateJobStatusStats{} // reset in case of retry
 		stmt, err := tx.PrepareContext(ctx, pq.CopyIn(ds.JobStatusTable, "job_id", "job_state", "attempt", "exec_time",
 			"retry_time", "error_code", "error_response", "parameters"))
 		if err != nil {
@@ -2362,18 +2429,29 @@ func (jd *Handle) updateJobStatusDSInTx(ctx context.Context, tx *Tx, ds dataSetT
 		defer func() { _ = stmt.Close() }()
 		for _, status := range statusList {
 			//  Handle the case when google analytics returns gif in response
-			if _, ok := updatedStates[status.WorkspaceId]; !ok {
-				updatedStates[status.WorkspaceId] = make(map[string]map[ParameterFilterT]struct{})
+			if _, ok := updatedStates[workspaceIDKey(status.WorkspaceId)]; !ok {
+				updatedStates[workspaceIDKey(status.WorkspaceId)] = make(map[jobStateKey]map[parameterFiltersKey]*UpdateJobStatusStats)
 			}
-			if _, ok := updatedStates[status.WorkspaceId][status.JobState]; !ok {
-				updatedStates[status.WorkspaceId][status.JobState] = make(map[ParameterFilterT]struct{})
+			if _, ok := updatedStates[workspaceIDKey(status.WorkspaceId)][jobStateKey(status.JobState)]; !ok {
+				updatedStates[workspaceIDKey(status.WorkspaceId)][jobStateKey(status.JobState)] = make(map[parameterFiltersKey]*UpdateJobStatusStats)
 			}
+			var parameters ParameterFilterList
+			var parametersKey parameterFiltersKey
 			if status.JobParameters != nil {
 				for _, param := range cacheParameterFilters {
 					v := gjson.GetBytes(status.JobParameters, param).Str
-					updatedStates[status.WorkspaceId][status.JobState][ParameterFilterT{Name: param, Value: v}] = struct{}{}
+					parameters = append(parameters, ParameterFilterT{Name: param, Value: v})
 				}
+				parametersKey = parameterFiltersKey(parameters.String())
+
+			}
+			pm, ok := updatedStates[workspaceIDKey(status.WorkspaceId)][jobStateKey(status.JobState)][parametersKey]
+			if !ok {
+				pm = &UpdateJobStatusStats{parameters: parameters}
+				updatedStates[workspaceIDKey(status.WorkspaceId)][jobStateKey(status.JobState)][parametersKey] = pm
 			}
+			pm.count++
+			pm.bytes += len(status.ErrorResponse)
 
 			if !utf8.ValidString(string(status.ErrorResponse)) {
 				status.ErrorResponse = []byte(`{}`)
@@ -2838,38 +2916,53 @@ func (jd *Handle) internalUpdateJobStatusInTx(ctx context.Context, tx *Tx, dsLis
 
 	tx.AddSuccessListener(func() {
 		// clear cache
-		for ds, dsKeys := range updatedStatesByDS {
-			if len(dsKeys) == 0 { // if no keys, we need to invalidate all keys
+		for ds, dsStats := range updatedStatesByDS {
+			if len(dsStats) == 0 { // if no keys, we need to invalidate all keys
 				jd.noResultsCache.Invalidate(ds.Index, "", nil, nil, nil)
 			}
-			for workspace, wsKeys := range dsKeys {
-				if len(wsKeys) == 0 { // if no keys, we need to invalidate all keys
-					jd.noResultsCache.Invalidate(ds.Index, workspace, nil, nil, nil)
+			for workspace, wsStats := range dsStats {
+				if len(wsStats) == 0 { // if no keys, we need to invalidate all keys
+					jd.noResultsCache.Invalidate(ds.Index, string(workspace), nil, nil, nil)
 				}
-				for state, parametersMap := range wsKeys {
-					stateList := []string{state}
-					if len(parametersMap) == 0 { // if no keys, we need to invalidate all keys
-						jd.noResultsCache.Invalidate(ds.Index, workspace, customValFilters, stateList, nil)
-					}
-					parameterFilters := lo.Keys(parametersMap)
-					jd.noResultsCache.Invalidate(ds.Index, workspace, customValFilters, stateList, parameterFilters)
+				for state, parametersStats := range wsStats {
+					stateList := []string{string(state)}
+					parameterFilters := lo.UniqBy( // gather unique parameter filters
+						lo.FlatMap(
+							lo.Values(parametersStats), // from all JobStatusMetrics
+							func(ujss *UpdateJobStatusStats, _ int) []ParameterFilterT {
+								return ujss.parameters
+							},
+						),
+						func(pf ParameterFilterT) string {
+							return pf.String() // uniqueness by string representation
+						},
+					)
+					// invalidate cache for this combination
+					jd.noResultsCache.Invalidate(ds.Index, string(workspace), customValFilters, stateList, parameterFilters)
 				}
 			}
 		}
 	})
+
+	// use the aggregated stats from updateJobStatusInTx
 	tx.AddSuccessListener(func() {
-		statTags := tags.getStatsTags(jd.tablePrefix)
-		statusCounters := make(map[string]lo.Tuple2[int, int], 0)
-		for i := range statusList {
-			t := statusCounters[statusList[i].JobState]
-			t.A++                                   // job count
-			t.B += len(statusList[i].ErrorResponse) // bytes count
-			statusCounters[statusList[i].JobState] = t
-		}
-		for state, count := range statusCounters {
-			statTags["jobState"] = state
-			jd.stats.NewTaggedStat("jobsdb_updated_jobs", stats.CountType, statTags).Count(count.A)
-			jd.stats.NewTaggedStat("jobsdb_updated_bytes", stats.CountType, statTags).Count(count.B)
+		merged := updateJobStatusStats{}
+		for _, dsStats := range updatedStatesByDS {
+			merged.Merge(dsStats)
+		}
+		statsByState := merged.StatsByState()
+		for state, parametersMap := range statsByState {
+			for _, metrics := range parametersMap {
+				statTags := tags.getStatsTags(jd.tablePrefix)
+				statTags["jobState"] = string(state)
+				if len(parameterFilters) == 0 { // only add parameter filters if not already set
+					for _, pf := range metrics.parameters {
+						statTags[pf.Name] = pf.Value
+					}
+				}
+				jd.stats.NewTaggedStat("jobsdb_updated_jobs", stats.CountType, statTags).Count(metrics.count)
+				jd.stats.NewTaggedStat("jobsdb_updated_bytes", stats.CountType, statTags).Count(metrics.bytes)
+			}
 		}
 	})
 
@@ -2881,7 +2974,7 @@ doUpdateJobStatusInTx updates the status of a batch of jobs
 customValFilters[] is passed, so we can efficiently mark empty cache
 Later we can move this to query
 */
-func (jd *Handle) doUpdateJobStatusInTx(ctx context.Context, tx *Tx, dsList []dataSetT, dsRangeList []dataSetRangeT, statusList []*JobStatusT, tags statTags) (updatedStatesByDS map[dataSetT]map[string]map[string]map[ParameterFilterT]struct{}, err error) {
+func (jd *Handle) doUpdateJobStatusInTx(ctx context.Context, tx *Tx, dsList []dataSetT, dsRangeList []dataSetRangeT, statusList []*JobStatusT, tags statTags) (updatedStatesByDS map[dataSetT]updateJobStatusStats, err error) {
 	if len(statusList) == 0 {
 		return updatedStatesByDS, err
 	}
@@ -2893,7 +2986,7 @@ func (jd *Handle) doUpdateJobStatusInTx(ctx context.Context, tx *Tx, dsList []da
 
 	// We scan through the list of jobs and map them to DS
 	var lastPos int
-	updatedStatesByDS = make(map[dataSetT]map[string]map[string]map[ParameterFilterT]struct{})
+	updatedStatesByDS = make(map[dataSetT]updateJobStatusStats)
 	for _, ds := range dsRangeList {
 		minID := ds.minJobID
 		maxID := ds.maxJobID
@@ -2914,7 +3007,7 @@ func (jd *Handle) doUpdateJobStatusInTx(ctx context.Context, tx *Tx, dsList []da
 						logger.NewIntField("prevPos", int64(i-1)),
 					)
 				}
-				var updatedStates map[string]map[string]map[ParameterFilterT]struct{}
+				var updatedStates updateJobStatusStats
 				updatedStates, err = jd.updateJobStatusDSInTx(ctx, tx, ds.ds, statusList[lastPos:i], tags)
 				if err != nil {
 					return updatedStatesByDS, err
@@ -2935,7 +3028,7 @@ func (jd *Handle) doUpdateJobStatusInTx(ctx context.Context, tx *Tx, dsList []da
 				logger.NewIntField("prevJobID", statusList[i-1].JobID),
 				logger.NewIntField("index", int64(i)),
 			)
-			var updatedStates map[string]map[string]map[ParameterFilterT]struct{}
+			var updatedStates updateJobStatusStats
 			updatedStates, err = jd.updateJobStatusDSInTx(ctx, tx, ds.ds, statusList[lastPos:i], tags)
 			if err != nil {
 				return updatedStatesByDS, err
@@ -2957,7 +3050,7 @@ func (jd *Handle) doUpdateJobStatusInTx(ctx context.Context, tx *Tx, dsList []da
 		jd.logger.Debugn("RangeEnd",
 			logger.NewIntField("jobID", statusList[lastPos].JobID),
 			logger.NewIntField("lenStatusList", int64(len(statusList))))
-		var updatedStates map[string]map[string]map[ParameterFilterT]struct{}
+		var updatedStates updateJobStatusStats
 		updatedStates, err = jd.updateJobStatusDSInTx(ctx, tx, dsList[len(dsList)-1], statusList[lastPos:], tags)
 		if err != nil {
 			return updatedStatesByDS, err