safe responsewriter usage in TryExec (#1490)

inside of TryExec we were writing directly to the response writer inside of a
goroutine, but TryExec can timeout and then get called again to a different
runner or even have the front end writing headers while TryExec is writing
headers.

one way to make this safe is to make a new response writer for TryExec to
write the response into, and only after the goroutine handling the response
has returned, from the TryExec goroutine we can copy the response back up as
the caller will not call TryExec again until it has returned (this is
seemingly part of the placer contract). unfortunately, we're already buffering
the response writer in the front end, too - it's possible we can get rid of
that but it may need further testing.

this adds an optimization when copying the request body from the LB to a
runner, since we're using request.GetBody() and returning a reader we
are familiar with that happens to just wrap a buffer's bytes (which we just
need multiple readers on, but the data doesn't change). anyway, this whole
interaction is unfortunate but kind of necessary due to needing to maneuver
into a protobuf, it seems like a worth it and somewhat ok abstraction wise
optimization.

additionally, this gets rid of passing the client response headers down into
the agent for detached functions. we don't need these since detached functions
are not responding with the functions response to the client, only a 202, this
was leading to races around writing the headers in retries too, but this is
just for posterity/correctness now.

updated the makefile/system test script so that I could run these faster to
repro, pretty handy, should add to other stuff too...

closes #1484

Author: rdallman

Makefile

@@ -70,9 +70,9 @@ test: checkfmt pull-images test-basic test-middleware test-extensions test-syste
 
 .PHONY: test-system
 test-system:
-	./system_test.sh sqlite3
-	./system_test.sh mysql
-	./system_test.sh postgres
+	./system_test.sh sqlite3 $(run)
+	./system_test.sh mysql $(run)
+	./system_test.sh postgres $(run)
 
 .PHONY: img-busybox
 img-busybox:

api/agent/lb_agent.go

@@ -28,14 +28,18 @@ type lbAgent struct {
 	callOpts      []CallOpt
 }
 
+// DetachedResponseWriter implements http.ResponseWriter without allowing
+// writes to the body or writing the headers from a call to Write or
+// WriteHeader, it is only intended to allow writing the status code in and
+// being able to fetch it later from Status()
 type DetachedResponseWriter struct {
-	Headers http.Header
+	headers http.Header
 	status  int
 	acked   chan struct{}
 }
 
 func (w *DetachedResponseWriter) Header() http.Header {
-	return w.Headers
+	return w.headers
 }
 
 func (w *DetachedResponseWriter) Write(data []byte) (int, error) {
@@ -51,9 +55,9 @@ func (w *DetachedResponseWriter) Status() int {
 	return w.status
 }
 
-func NewDetachedResponseWriter(h http.Header, statusCode int) *DetachedResponseWriter {
+func NewDetachedResponseWriter(statusCode int) *DetachedResponseWriter {
 	return &DetachedResponseWriter{
-		Headers: h,
+		headers: make(http.Header),
 		status:  statusCode,
 		acked:   make(chan struct{}, 1),
 	}

api/agent/runner_client.go

@@ -1,6 +1,7 @@
 package agent
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"encoding/hex"
@@ -161,6 +162,76 @@ func (r *gRPCRunner) Status(ctx context.Context) (*pool.RunnerStatus, error) {
 
 // implements Runner
 func (r *gRPCRunner) TryExec(ctx context.Context, call pool.RunnerCall) (bool, error) {
+
+	// we need to get a response writer that is safe for us to use. the only reason we need this,
+	// ostensibly, is because io operations can't be timed out so we need to do receiveFromRunner
+	// in a goroutine, and we may return from here from a timeout while recv is running, leading
+	// to a race when a placer calls TryExec with the same ResponseWriter we have here. blech.
+	respBuffer := bufPool.Get().(*bytes.Buffer)
+	respBuffer.Reset()
+	defer func() {
+		if ctx.Err() == nil { // this is only safe if we don't time out (receiveFromRunner returned/not called)
+			bufPool.Put(respBuffer)
+		}
+	}()
+
+	writer := syncResponseWriter{
+		headers: make(http.Header),
+		Buffer:  respBuffer,
+	}
+
+	safeCall := callWithResponseWriter(call, &writer)
+	placed, err := r.tryExec(ctx, safeCall)
+	if err != nil || !placed {
+		return placed, err
+	}
+
+	// now we can write to the actual response writer from this thread (so long
+	// as the caller waits, we're playing jenga ofc)
+	rw := call.ResponseWriter()
+	copyHeaders(rw.Header(), writer.Header())
+	rw.WriteHeader(writer.status)
+	io.Copy(rw, writer) // TODO(reed): this is also a buffer->buffer operation :( but it means no errors
+	return true, nil
+}
+
+func callWithResponseWriter(call pool.RunnerCall, rw http.ResponseWriter) pool.RunnerCall {
+	return &wrapCall{call, rw}
+}
+
+// wrapCall implements pool.RunnerCall but bypasses the embedded RunnerCall's ResponseWriter() method
+// TODO this is the worst thing we've tried, except for all the other things we've tried.
+type wrapCall struct {
+	pool.RunnerCall
+	rw http.ResponseWriter
+}
+
+func (c *wrapCall) ResponseWriter() http.ResponseWriter {
+	return c.rw
+}
+
+func copyHeaders(dst, src http.Header) {
+	for k, vs := range src {
+		for _, v := range vs {
+			dst.Add(k, v)
+		}
+	}
+}
+
+// TODO(reed): this is copied. and not only does it make sense in both places for different
+// reasons, it makes sense in another place too. need to reconsider ifaces
+type syncResponseWriter struct {
+	headers http.Header
+	status  int
+	*bytes.Buffer
+}
+
+var _ http.ResponseWriter = new(syncResponseWriter) // nice compiler errors
+
+func (s *syncResponseWriter) Header() http.Header  { return s.headers }
+func (s *syncResponseWriter) WriteHeader(code int) { s.status = code }
+
+func (r *gRPCRunner) tryExec(ctx context.Context, call pool.RunnerCall) (bool, error) {
 	log := common.Logger(ctx).WithField("runner_addr", r.address)
 
 	log.Debug("Attempting to place call")
@@ -342,6 +413,7 @@ func recordFinishStats(ctx context.Context, msg *pb.CallFinished, c pool.RunnerC
 		statsLBAgentRunnerSchedLatency(ctx, runnerSchedLatency)
 		statsLBAgentRunnerExecLatency(ctx, runnerExecLatency)
 
+		// TODO: this is not safe to be called from within receiveFromRunner, it may get called each retry (data race, but also incorrect)
 		c.AddUserExecutionTime(runnerExecLatency)
 	}
 }

api/server/runner_fninvoke.go

@@ -96,7 +96,7 @@ func (s *Server) fnInvoke(resp http.ResponseWriter, req *http.Request, app *mode
 
 	isDetached := req.Header.Get("Fn-Invoke-Type") == models.TypeDetached
 	if isDetached {
-		writer = agent.NewDetachedResponseWriter(resp.Header(), 202)
+		writer = agent.NewDetachedResponseWriter(202)
 	} else {
 		writer = &syncResponseWriter{
 			headers: resp.Header(),

system_test.sh

@@ -6,6 +6,7 @@ source ./helpers.sh
 remove_containers ${CONTEXT}
 
 DB_NAME=$1
+shift # later usage
 export FN_DB_URL=$(spawn_${DB_NAME} ${CONTEXT})
 
 # avoid port conflicts with api_test.sh which are run in parallel
@@ -23,8 +24,15 @@ export FN_LOG_LEVEL=debug
 #
 export SYSTEM_TEST_PROMETHEUS_FILE=./prometheus.${DB_NAME}.txt
 
+run="$@"
+
+if [ ! -z "$run" ]
+then
+  run="-run $run"
+fi
+
 cd test/fn-system-tests
-go test -v ./...
+go test $run -v ./...
 cd ../../
 
 remove_containers ${CONTEXT}