Initial Refactor (#1234)

* Inital Refactor

Removing the repeated logic exposed some problems with the reponse
writers.

Currently, the trigger writer was overlaid on part of the header
writing. The main invoke blog writing into the different levels of the
overlays at different points in the logic.

Instead, by extending the types and embedded structs, the writer is
more transparent. So, at the end of the flow it goes over all the
headers available and removes our prefixes. This lets the invoke logic
just write to the top level.

Going to continue after lunch to try and remove some of the layers and
param passing.

* Try and repeat concurrency failure

* Nested FromHTTPFnRequest inside FromHTTPTriggerRequest

* Consolidate buffer pooling logic

* go fmt yourself

* fix import

api/server/runner_fninvoke.go

@@ -20,7 +20,15 @@ var (
 	bufPool = &sync.Pool{New: func() interface{} { return new(bytes.Buffer) }}
 )
 
-var _ http.ResponseWriter = new(syncResponseWriter)
+type ResponseBufferingWriter interface {
+	http.ResponseWriter
+	io.Reader
+	Status() int
+	GetBuffer() *bytes.Buffer
+	SetBuffer(*bytes.Buffer)
+}
+
+var _ ResponseBufferingWriter = new(syncResponseWriter)
 
 // implements http.ResponseWriter
 // this little guy buffers responses from user containers and lets them still
@@ -33,8 +41,13 @@ type syncResponseWriter struct {
 	*bytes.Buffer
 }
 
-func (s *syncResponseWriter) Header() http.Header  { return s.headers }
-func (s *syncResponseWriter) WriteHeader(code int) { s.status = code }
+func (s *syncResponseWriter) Header() http.Header { return s.headers }
+
+// By storing the status here, we effectively buffer the response
+func (s *syncResponseWriter) WriteHeader(code int)        { s.status = code }
+func (s *syncResponseWriter) Status() int                 { return s.status }
+func (s *syncResponseWriter) GetBuffer() *bytes.Buffer    { return s.Buffer }
+func (s *syncResponseWriter) SetBuffer(buf *bytes.Buffer) { s.Buffer = buf }
 
 // handleFnInvokeCall executes the function, for router handlers
 func (s *Server) handleFnInvokeCall(c *gin.Context) {
@@ -66,63 +79,71 @@ func (s *Server) handleFnInvokeCall2(c *gin.Context) error {
 }
 
 func (s *Server) ServeFnInvoke(c *gin.Context, app *models.App, fn *models.Fn) error {
-	// TODO: we should combine this logic with trigger, which just wraps this block with some headers wizardry
-	// TODO: we should get rid of the buffers, and stream back (saves memory (+splice), faster (splice), allows streaming, don't have to cap resp size)
-	buf := bufPool.Get().(*bytes.Buffer)
-	buf.Reset()
-	writer := syncResponseWriter{
-		Buffer:  buf,
+	writer := &syncResponseWriter{
 		headers: c.Writer.Header(),
 	}
-	defer bufPool.Put(buf) // TODO need to ensure this is safe with Dispatch?
 
-	call, err := s.agent.GetCall(
-		agent.WithWriter(&writer), // XXX (reed): order matters [for now]
-		agent.FromHTTPFnRequest(app, fn, c.Request),
-	)
+	call, err := s.agent.GetCall(agent.WithWriter(writer), // XXX (reed): order matters [for now]
+		agent.FromHTTPFnRequest(app, fn, c.Request))
+
 	if err != nil {
 		return err
 	}
 
+	return s.fnInvoke(c, app, fn, writer, call)
+}
+
+func (s *Server) fnInvoke(c *gin.Context, app *models.App, fn *models.Fn, writer ResponseBufferingWriter, call agent.Call) error {
+	// TODO: we should get rid of the buffers, and stream back (saves memory (+splice), faster (splice), allows streaming, don't have to cap resp size)
+	buf := bufPool.Get().(*bytes.Buffer)
+	buf.Reset()
+	var submitErr error
+	defer func() {
+		if buf.Len() == 0 && submitErr == nil {
+			bufPool.Put(buf) // TODO need to ensure this is safe with Dispatch?
+		}
+	}()
+	writer.SetBuffer(buf)
+
 	model := call.Model()
 	{ // scope this, to disallow ctx use outside of this scope. add id for handleV1ErrorResponse logger
 		ctx, _ := common.LoggerWithFields(c.Request.Context(), logrus.Fields{"id": model.ID})
 		c.Request = c.Request.WithContext(ctx)
 	}
 
-	err = s.agent.Submit(call)
-	if err != nil {
+	submitErr = s.agent.Submit(call)
+	if submitErr != nil {
 		// NOTE if they cancel the request then it will stop the call (kind of cool),
 		// we could filter that error out here too as right now it yells a little
-		if err == models.ErrCallTimeoutServerBusy || err == models.ErrCallTimeout {
+		if submitErr == models.ErrCallTimeoutServerBusy || submitErr == models.ErrCallTimeout {
 			// TODO maneuver
 			// add this, since it means that start may not have been called [and it's relevant]
 			c.Writer.Header().Add("XXX-FXLB-WAIT", time.Now().Sub(time.Time(model.CreatedAt)).String())
 		}
-		return err
+		return submitErr
 	}
-
 	// if they don't set a content-type - detect it
 	// TODO: remove this after removing all the formats (too many tests to scrub til then)
 	if writer.Header().Get("Content-Type") == "" {
 		// see http.DetectContentType, the go server is supposed to do this for us but doesn't appear to?
 		var contentType string
 		jsonPrefix := [1]byte{'{'} // stack allocated
-		if bytes.HasPrefix(buf.Bytes(), jsonPrefix[:]) {
+		if bytes.HasPrefix(writer.GetBuffer().Bytes(), jsonPrefix[:]) {
 			// try to detect json, since DetectContentType isn't a hipster.
 			contentType = "application/json; charset=utf-8"
 		} else {
-			contentType = http.DetectContentType(buf.Bytes())
+			contentType = http.DetectContentType(writer.GetBuffer().Bytes())
 		}
 		writer.Header().Set("Content-Type", contentType)
 	}
 
-	writer.Header().Set("Content-Length", strconv.Itoa(int(buf.Len())))
+	writer.Header().Set("Content-Length", strconv.Itoa(int(writer.GetBuffer().Len())))
 
-	if writer.status > 0 {
-		c.Writer.WriteHeader(writer.status)
+	if writer.Status() > 0 {
+		c.Writer.WriteHeader(writer.Status())
 	}
-	io.Copy(c.Writer, &writer)
+
+	io.Copy(c.Writer, writer)
 
 	return nil
 }