fix: fix-up context handling (#305)

examples/301-http-pipeline.py

@@ -13,14 +13,15 @@
 * Due to the lack of state, do not use HTTP sources and sinks to connect steps. 
 """)
      .annotate("dataflow.argoproj.io/needs", "dataflow-103-http-main-source-default-secret.yaml")
-     .step(
-        (cron('*/3 * * * * *')
-         .cat('cron')
-         .http('https://http-main/sources/default', insecureSkipVerify=True, headers=[{'name': "Authorization", "value": "Bearer my-bearer-token"}])
-         ))
      .step(
         (http(serviceName='http-main')
          .cat('main')
          .log()
          ))
+     .step(
+        (cron('*/3 * * * * *')
+         .cat('cron')
+         .http('https://http-main/sources/default', insecureSkipVerify=True,
+               headers=[{'name': "Authorization", "value": "Bearer my-bearer-token"}])
+         ))
      .save())

examples/301-http-pipeline.yaml

@@ -13,6 +13,13 @@ metadata:
   name: 301-http
 spec:
   steps:
+  - cat: {}
+    name: main
+    sinks:
+    - log: {}
+    sources:
+    - http:
+        serviceName: http-main
   - cat: {}
     name: cron
     sinks:
@@ -25,10 +32,3 @@ spec:
     sources:
     - cron:
         schedule: '*/3 * * * * *'
-  - cat: {}
-    name: main
-    sinks:
-    - log: {}
-    sources:
-    - http:
-        serviceName: http-main

runner/sidecar/in.go

@@ -15,7 +15,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
-func connectIn(ctx context.Context, sink func([]byte) error) (func(context.Context, []byte) error, error) {
+func connectIn(ctx context.Context, sink func(context.Context, []byte) error) (func(context.Context, []byte) error, error) {
 	inFlight := promauto.NewGauge(prometheus.GaugeOpts{
 		Subsystem:   "input",
 		Name:        "inflight",
@@ -81,7 +81,7 @@ func connectIn(ctx context.Context, sink func([]byte) error) (func(context.Conte
 					return fmt.Errorf("failed to send to main: %q %q", resp.Status, body)
 				}
 				if resp.StatusCode == 201 {
-					return sink(body)
+					return sink(ctx, body)
 				}
 			}
 			return nil

runner/sidecar/out.go

@@ -13,12 +13,12 @@ import (
 	runtimeutil "k8s.io/apimachinery/pkg/util/runtime"
 )
 
-func connectOut(toSinks func([]byte) error) {
-	connectOutFIFO(toSinks)
-	connectOutHTTP(toSinks)
+func connectOut(ctx context.Context, sink func(context.Context, []byte) error) {
+	connectOutFIFO(ctx, sink)
+	connectOutHTTP(sink)
 }
 
-func connectOutHTTP(f func([]byte) error) {
+func connectOutHTTP(sink func(context.Context, []byte) error) {
 	logger.Info("HTTP out interface configured")
 	v, err := ioutil.ReadFile(dfv1.PathAuthorization)
 	if err != nil {
@@ -37,7 +37,7 @@ func connectOutHTTP(f func([]byte) error) {
 			_, _ = w.Write([]byte(err.Error()))
 			return
 		}
-		if err := f(data); err != nil {
+		if err := sink(r.Context(), data); err != nil {
 			logger.Error(err, "failed to send message from main to sink")
 			w.WriteHeader(500)
 			_, _ = w.Write([]byte(err.Error()))
@@ -47,7 +47,7 @@ func connectOutHTTP(f func([]byte) error) {
 	})
 }
 
-func connectOutFIFO(f func([]byte) error) {
+func connectOutFIFO(ctx context.Context, sink func(context.Context, []byte) error) {
 	logger.Info("FIFO out interface configured")
 	go func() {
 		defer runtimeutil.HandleCrash()
@@ -63,7 +63,7 @@ func connectOutFIFO(f func([]byte) error) {
 			logger.Info("opened output FIFO")
 			scanner := bufio.NewScanner(fifo)
 			for scanner.Scan() {
-				if err := f(scanner.Bytes()); err != nil {
+				if err := sink(ctx, scanner.Bytes()); err != nil {
 					return fmt.Errorf("failed to send message from main to sink: %w", err)
 				}
 			}

runner/sidecar/sidecar.go

@@ -105,7 +105,7 @@ func Exec(ctx context.Context) error {
 
 	addStopHook(patchStepStatusHook)
 
-	toSinks, err := connectSinks(ctx)
+	sink, err := connectSinks(ctx)
 	if err != nil {
 		return err
 	}
@@ -146,7 +146,7 @@ func Exec(ctx context.Context) error {
 		w.WriteHeader(204)
 	})
 
-	connectOut(toSinks)
+	connectOut(ctx, sink)
 
 	server := &http.Server{Addr: "localhost:3569"}
 	addStopHook(func(ctx context.Context) error {
@@ -175,12 +175,12 @@ func Exec(ctx context.Context) error {
 		logger.Info("HTTPS server shutdown")
 	}()
 
-	toMain, err := connectIn(ctx, toSinks)
+	process, err := connectIn(ctx, sink)
 	if err != nil {
 		return err
 	}
 
-	if err := connectSources(ctx, toMain); err != nil {
+	if err := connectSources(ctx, process); err != nil {
 		return err
 	}
 

runner/sidecar/sink/db/db.go

@@ -64,7 +64,7 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, x dfv1.DBS
 	}, nil
 }
 
-func (d dbSink) Sink(msg []byte) error {
+func (d dbSink) Sink(ctx context.Context, msg []byte) error {
 	tx, err := d.db.Begin()
 	if err != nil {
 		return fmt.Errorf("failed to start a transaction: %w", err)

runner/sidecar/sink/http/http.go

@@ -49,8 +49,8 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, x dfv1.HTT
 	}, nil
 }
 
-func (h httpSink) Sink(msg []byte) error {
-	req, err := http.NewRequest("POST", h.url, bytes.NewBuffer(msg))
+func (h httpSink) Sink(ctx context.Context, msg []byte) error {
+	req, err := http.NewRequestWithContext(ctx, "POST", h.url, bytes.NewBuffer(msg))
 	if err != nil {
 		return fmt.Errorf("failed to create HTTP request: %w", err)
 	}

runner/sidecar/sink/interface.go

@@ -1,5 +1,7 @@
 package sink
 
+import "context"
+
 type Interface interface {
-	Sink(msg []byte) error
+	Sink(ctx context.Context, msg []byte) error
 }

runner/sidecar/sink/kafka/kafka.go

@@ -19,20 +19,20 @@ import (
 var logger = sharedutil.NewLogger()
 
 type producer interface {
-	SendMessage(msg *sarama.ProducerMessage) error
+	SendMessage(ctx context.Context, msg *sarama.ProducerMessage) error
 	io.Closer
 }
 
 type asyncProducer struct{ sarama.AsyncProducer }
 
-func (s asyncProducer) SendMessage(msg *sarama.ProducerMessage) error {
+func (s asyncProducer) SendMessage(ctx context.Context, msg *sarama.ProducerMessage) error {
 	s.AsyncProducer.Input() <- msg
 	return nil
 }
 
 type syncProducer struct{ sarama.SyncProducer }
 
-func (s syncProducer) SendMessage(msg *sarama.ProducerMessage) error {
+func (s syncProducer) SendMessage(ctx context.Context, msg *sarama.ProducerMessage) error {
 	_, _, err := s.SyncProducer.SendMessage(msg)
 	return err
 }
@@ -103,8 +103,8 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, x dfv1.Kaf
 	}
 }
 
-func (h kafkaSink) Sink(msg []byte) error {
-	return h.producer.SendMessage(&sarama.ProducerMessage{Value: sarama.ByteEncoder(msg), Topic: h.topic})
+func (h kafkaSink) Sink(ctx context.Context, msg []byte) error {
+	return h.producer.SendMessage(ctx, &sarama.ProducerMessage{Value: sarama.ByteEncoder(msg), Topic: h.topic})
 }
 
 func (h kafkaSink) Close() error {

runner/sidecar/sink/log/log.go

@@ -1,6 +1,8 @@
 package logsink
 
 import (
+	"context"
+
 	dfv1 "github.com/argoproj-labs/argo-dataflow/api/v1alpha1"
 	"github.com/argoproj-labs/argo-dataflow/runner/sidecar/sink"
 	sharedutil "github.com/argoproj-labs/argo-dataflow/shared/util"
@@ -16,7 +18,7 @@ func New(x dfv1.Log) sink.Interface {
 	return logSink{truncate: x.Truncate}
 }
 
-func (s logSink) Sink(msg []byte) error {
+func (s logSink) Sink(ctx context.Context, msg []byte) error {
 	text := string(msg)
 	if s.truncate != nil && len(text) > int(*s.truncate) {
 		text = text[0:*s.truncate]

runner/sidecar/sink/s3/s3.go

@@ -73,7 +73,7 @@ func New(ctx context.Context, secretInterface v1.SecretInterface, x dfv1.S3Sink)
 	return s3Sink{client: s3.New(options), bucket: x.Bucket}, nil
 }
 
-func (h s3Sink) Sink(msg []byte) error {
+func (h s3Sink) Sink(ctx context.Context, msg []byte) error {
 	m := &message{}
 	if err := json.Unmarshal(msg, m); err != nil {
 		return err
@@ -82,7 +82,7 @@ func (h s3Sink) Sink(msg []byte) error {
 	if err != nil {
 		return fmt.Errorf("failed to open %q: %w", m.Path, err)
 	}
-	_, err = h.client.PutObject(context.Background(), &s3.PutObjectInput{
+	_, err = h.client.PutObject(ctx, &s3.PutObjectInput{
 		Bucket: &h.bucket,
 		Key:    &m.Key,
 		Body:   f,

runner/sidecar/sink/stan/stan.go

@@ -70,7 +70,7 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, namespace,
 	}, nil
 }
 
-func (s stanSink) Sink(msg []byte) error {
+func (s stanSink) Sink(ctx context.Context, msg []byte) error {
 	return s.conn.Publish(s.subject, msg)
 }
 

runner/sidecar/sink/volume/volume.go

@@ -1,6 +1,8 @@
 package volume
 
 import (
+	"context"
+
 	"github.com/argoproj-labs/argo-dataflow/runner/sidecar/sink"
 )
 
@@ -10,6 +12,6 @@ func New() (sink.Interface, error) {
 	return volumeSink{}, nil
 }
 
-func (h volumeSink) Sink([]byte) error {
+func (h volumeSink) Sink(ctx context.Context, msg []byte) error {
 	return nil
 }

runner/sidecar/sinks.go

@@ -19,7 +19,7 @@ import (
 	"github.com/paulbellamy/ratecounter"
 )
 
-func connectSinks(ctx context.Context) (func([]byte) error, error) {
+func connectSinks(ctx context.Context) (func(context.Context, []byte) error, error) {
 	sinks := map[string]sink.Interface{}
 	rateCounters := map[string]*ratecounter.RateCounter{}
 	for _, sink := range step.Spec.Sinks {
@@ -79,14 +79,14 @@ func connectSinks(ctx context.Context) (func([]byte) error, error) {
 		}
 	}
 
-	return func(msg []byte) error {
+	return func(ctx context.Context, msg []byte) error {
 		for sinkName, f := range sinks {
 			counter := rateCounters[sinkName]
 			counter.Incr(1)
 			withLock(func() {
 				step.Status.SinkStatues.IncrTotal(sinkName, replica, rateToResourceQuantity(counter), uint64(len(msg)))
 			})
-			if err := f.Sink(msg); err != nil {
+			if err := f.Sink(ctx, msg); err != nil {
 				withLock(func() { step.Status.SinkStatues.IncrErrors(sinkName, replica) })
 				return err
 			}

runner/sidecar/source/cron/cron.go

@@ -19,7 +19,7 @@ type cronSource struct {
 	crn *cron.Cron
 }
 
-func New(x dfv1.Cron, f source.Func) (source.Interface, error) {
+func New(ctx context.Context, x dfv1.Cron, process source.Process) (source.Interface, error) {
 	crn := cron.New(
 		cron.WithParser(cron.NewParser(cron.SecondOptional|cron.Minute|cron.Hour|cron.Dom|cron.Month|cron.Dow|cron.Descriptor)),
 		cron.WithChain(cron.Recover(logger)),
@@ -32,7 +32,7 @@ func New(x dfv1.Cron, f source.Func) (source.Interface, error) {
 
 	_, err := crn.AddFunc(x.Schedule, func() {
 		msg := []byte(time.Now().Format(x.Layout))
-		if err := f(context.Background(), msg); err != nil {
+		if err := process(ctx, msg); err != nil {
 			logger.Error(err, "failed to process message")
 		}
 	})

runner/sidecar/source/db/db.go

@@ -30,7 +30,7 @@ type dbSource struct {
 	db *sql.DB
 }
 
-func New(ctx context.Context, secretInterface corev1.SecretInterface, clusterName, namespace, pipelineName, stepName string, replica int, sourceName string, x dfv1.DBSource, f source.Func) (source.Interface, error) {
+func New(ctx context.Context, secretInterface corev1.SecretInterface, clusterName, namespace, pipelineName, stepName, sourceName string, x dfv1.DBSource, process source.Process) (source.Interface, error) {
 	dataSource, err := getDataSource(ctx, secretInterface, x)
 	if err != nil {
 		return nil, fmt.Errorf("failed to find data source: %w", err)
@@ -74,7 +74,7 @@ func New(ctx context.Context, secretInterface corev1.SecretInterface, clusterNam
 					if err != nil {
 						return fmt.Errorf("failed to marshal to json: %w", err)
 					}
-					if err := f(ctx, jsonData); err != nil {
+					if err := process(ctx, jsonData); err != nil {
 						return fmt.Errorf("failed to process data: %w", err)
 					}
 					offset = fmt.Sprintf("%v", d[x.OffsetColumn])

runner/sidecar/source/http/http.go

@@ -1,7 +1,6 @@
 package http
 
 import (
-	"context"
 	"io/ioutil"
 	"net/http"
 
@@ -12,7 +11,7 @@ type httpSource struct {
 	ready bool
 }
 
-func New(sourceName, authorization string, f source.Func) source.Interface {
+func New(sourceName, authorization string, process source.Process) source.Interface {
 	h := &httpSource{ready: true}
 	http.HandleFunc("/sources/"+sourceName, func(w http.ResponseWriter, r *http.Request) {
 		if r.Header.Get("Authorization") != authorization {
@@ -30,7 +29,7 @@ func New(sourceName, authorization string, f source.Func) source.Interface {
 			_, _ = w.Write([]byte(err.Error()))
 			return
 		}
-		if err := f(context.Background(), msg); err != nil {
+		if err := process(r.Context(), msg); err != nil {
 			w.WriteHeader(500)
 			_, _ = w.Write([]byte(err.Error()))
 		} else {

runner/sidecar/source/kafka/handler.go

@@ -7,7 +7,7 @@ import (
 )
 
 type handler struct {
-	f            source.Func
+	process      source.Process
 	i            int
 	manualCommit bool
 }
@@ -25,14 +25,14 @@ func (handler) Cleanup(_ sarama.ConsumerGroupSession) error {
 func (h handler) ConsumeClaim(sess sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
 	logger.Info("starting consuming claim", "partition", claim.Partition())
 	for msg := range claim.Messages() {
-		if err := h.f(sess.Context(), msg.Value); err != nil {
+		if err := h.process(sess.Context(), msg.Value); err != nil {
 			logger.Error(err, "failed to process message")
 		} else {
 			sess.MarkMessage(msg, "")
-		}
-		h.i++
-		if h.manualCommit && h.i%dfv1.CommitN == 0 {
-			sess.Commit()
+			h.i++
+			if h.manualCommit && h.i%dfv1.CommitN == 0 {
+				sess.Commit()
+			}
 		}
 	}
 	return nil