fix(go/exporter): Make sure all channels readers/writers are context-cancellable (google#4223)

When the context was being cancelled, the `ecosystemRouter` seems to
have been getting stuck, causing go to ungracefully terminate the
program due to deadlock.

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: michaelkedar

go/cmd/exporter/exporter.go

@@ -138,7 +138,18 @@ func ecosystemRouter(ctx context.Context, inCh <-chan *osvschema.Vulnerability,
 
 	allEcosystemWorker := newAllEcosystemWorker(ctx, outCh, &workersWg)
 
-	for vuln := range inCh {
+RouterLoop:
+	for {
+		var vuln *osvschema.Vulnerability
+		var ok bool
+		select {
+		case <-ctx.Done():
+			break RouterLoop
+		case vuln, ok = <-inCh:
+			if !ok {
+				break RouterLoop
+			}
+		}
 		vulnCounter++
 		ecosystems := make(map[string]struct{})
 		for _, aff := range vuln.GetAffected() {
@@ -164,15 +175,27 @@ func ecosystemRouter(ctx context.Context, inCh <-chan *osvschema.Vulnerability,
 				worker = newEcosystemWorker(ctx, eco, outCh, &workersWg)
 				workers[eco] = worker
 			}
-			worker.inCh <- vuln
+			select {
+			case worker.inCh <- vuln:
+			case <-ctx.Done():
+				break RouterLoop
+			}
+		}
+		select {
+		case allEcosystemWorker.inCh <- vulnAndEcos{Vulnerability: vuln, ecosystems: ecoNames}:
+		case <-ctx.Done():
+			break RouterLoop
 		}
-		allEcosystemWorker.inCh <- vulnAndEcos{Vulnerability: vuln, ecosystems: ecoNames}
 	}
 
 	for _, worker := range workers {
 		worker.Finish()
 	}
 	allEcosystemWorker.Finish()
 	workersWg.Wait()
-	logger.Info("ecosystem router finished, all vulnerabilities dispatched", slog.Int("total_vulnerabilities", vulnCounter))
+	if ctx.Err() == nil {
+		logger.Info("ecosystem router finished, all vulnerabilities dispatched", slog.Int("total_vulnerabilities", vulnCounter))
+	} else {
+		logger.Info("ecosystem router cancelled", slog.Any("err", ctx.Err()))
+	}
 }