[bugfix] Improve handling of context cancel in DBWorkManager to be more responsive

pkg/goDB/DBWorkManager.go

@@ -473,38 +473,37 @@ func (w *DBWorkManager) grabAndProcessWorkload(ctx context.Context, wg *sync.Wai
 			}
 		}()
 
-		var workload DBWorkload
-		for chanOpen := true; chanOpen; {
-			select {
-			case <-ctx.Done():
-				// query was cancelled, exit
-				logger.Infof("Query cancelled (workload %d / %d)...", w.nWorkloadsProcessed.Load(), w.nWorkloads)
-				return
-			case workload, chanOpen = <-workloadChan:
-				if chanOpen {
-					resultMap := hashmap.NewAggFlowMapWithMetadata()
-					for _, workDir := range workload.workDirs {
-
-						if memPool != nil {
-							workDir.SetMemPool(memPool)
-						}
-
-						// if there is an error during one of the read jobs, throw a syslog message and terminate
-						err := w.readBlocksAndEvaluate(workDir, enc, &resultMap)
-						if err != nil {
-							logger.Error(err)
-							mapChan <- hashmap.NilAggFlowMapWithMetadata
-							return
-						}
+		for workload := range workloadChan {
+			resultMap := hashmap.NewAggFlowMapWithMetadata()
+			for _, workDir := range workload.workDirs {
+				select {
+				case <-ctx.Done():
+
+					// query was cancelled, exit
+					logger.Infof("Query cancelled (workload %d / %d)...", w.nWorkloadsProcessed.Load(), w.nWorkloads)
+					return
+				default:
+
+					// check if a memory pool is available
+					if memPool != nil {
+						workDir.SetMemPool(memPool)
 					}
 
-					// Workload is counted, but we only add it to the final result if we got any entries
-					w.nWorkloadsProcessed.Add(1)
-					if resultMap.Len() > 0 {
-						mapChan <- resultMap
+					// if there is an error during one of the read jobs, throw a syslog message and terminate
+					err := w.readBlocksAndEvaluate(workDir, enc, &resultMap)
+					if err != nil {
+						logger.Error(err)
+						mapChan <- hashmap.NilAggFlowMapWithMetadata
+						return
 					}
 				}
 			}
+
+			// Workload is counted, but we only add it to the final result if we got any entries
+			w.nWorkloadsProcessed.Add(1)
+			if resultMap.Len() > 0 {
+				mapChan <- resultMap
+			}
 		}
 	}()
 }

pkg/goDB/engine/query.go

@@ -198,7 +198,7 @@ func (qr *QueryRunner) RunStatement(ctx context.Context, stmt *query.Statement)
 	result.Summary.Last = tSpanLast
 
 	// If enabled, run a live query in the background / parallel to the DB query and put the results on the same output channel
-	liveQueryWG := qr.runLiveQuery(ctx, mapChan, stmt)
+	liveQueryWG := qr.runLiveQuery(queryCtx, mapChan, stmt)
 
 	// spawn reader processing units and make them work on the individual DB blocks
 	// processing by interface is sequential, e.g. for multi-interface queries