Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SkipResponseWriter #3537

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions codegen/service/templates/service.go.tpl
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@
type Service interface {
{{- range .Methods }}
{{ comment .Description }}
{{- if .SkipResponseBodyEncodeDecode }}
{{ comment "\nIf body implements [io.WriterTo], that implementation will be used instead. Consider [goa.design/goa/v3/pkg.SkipResponseWriter] to adapt existing implementations." }}
{{- end }}
{{- if .ViewedResult }}
{{- if not .ViewedResult.ViewName }}
{{ comment "The \"view\" return value must have one of the following views" }}
Expand Down
16 changes: 16 additions & 0 deletions http/codegen/templates/server_handler_init.go.tpl
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,22 @@ func {{ .HandlerInit }}(
{{- if .Method.SkipResponseBodyEncodeDecode }}
o := res.(*{{ .ServicePkgName }}.{{ .Method.ResponseStruct }})
defer o.Body.Close()
if wt, ok := o.Body.(io.WriterTo); ok {
n, err := wt.WriteTo(w)
if err != nil {
if n == 0 {
if err := encodeError(ctx, w, err); err != nil {
errhandler(ctx, w, err)
}
} else {
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
panic(http.ErrAbortHandler) // too late to write an error
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if it might not be better to call errhandler here instead of panicking. Typical implementations notify the observability stack of the problem (log, error span, event etc.).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was a bit torn there. Decided to keep the semantics with the non-io.WriterTo version, which also panics. Should probably change both places if it's going to change, but this also feels like a nice bit of code to pull out into a helper for the generated code to call.

		if wt, ok := o.Body.(io.WriterTo); ok {
			n, err := wt.WriteTo(w)
			if err != nil {
				if n == 0 {
					if err := encodeError(ctx, w, err); err != nil {
						errhandler(ctx, w, err)
					}
				} else {
					if f, ok := w.(http.Flusher); ok {
						f.Flush()
					}
					panic(http.ErrAbortHandler) // too late to write an error
				}
			}
			return
		}
		// handle immediate read error like a returned error
		buf := bufio.NewReader(o.Body)
		if _, err := buf.Peek(1); err != nil && err != io.EOF {
			if err := encodeError(ctx, w, err); err != nil {
				errhandler(ctx, w, err)
			}
			return
		}
		if err := encodeResponse(ctx, w, o.Result); err != nil {
			errhandler(ctx, w, err)
			return
		}
		if _, err := io.Copy(w, buf); err != nil {
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			panic(http.ErrAbortHandler) // too late to write an error
		}

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah good point, hmm might as well stay consistent in both cases and stick with panic then - thanks for bringing this up!

}
}
return
}
// handle immediate read error like a returned error
buf := bufio.NewReader(o.Body)
if _, err := buf.Peek(1); err != nil && err != io.EOF {
Expand Down
16 changes: 16 additions & 0 deletions http/codegen/testdata/handler_init_functions.go
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,22 @@ func NewMethodSkipResponseBodyEncodeDecodeHandler(
}
o := res.(*serviceskipresponsebodyencodedecode.MethodSkipResponseBodyEncodeDecodeResponseData)
defer o.Body.Close()
if wt, ok := o.Body.(io.WriterTo); ok {
n, err := wt.WriteTo(w)
if err != nil {
if n == 0 {
if err := encodeError(ctx, w, err); err != nil {
errhandler(ctx, w, err)
}
} else {
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
panic(http.ErrAbortHandler) // too late to write an error
}
}
return
}
// handle immediate read error like a returned error
buf := bufio.NewReader(o.Body)
if _, err := buf.Peek(1); err != nil && err != io.EOF {
Expand Down
66 changes: 66 additions & 0 deletions pkg/skip_response_writer.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
package goa

import (
"io"
"sync"
"sync/atomic"
)

// SkipResponseWriter converts an io.WriterTo into a io.ReadCloser.
// The Read/Close methods this function returns will pipe the Write calls that wt makes, to implement a Reader that has the written bytes.
// If Read is called Close must also be called to avoid leaking memory.
// The returned value implements io.WriterTo as well, so the generated handler will call that instead of the Read method.
//
// Server handlers that use SkipResponseBodyEncodeDecode() io.ReadCloser as a return type.
func SkipResponseWriter(wt io.WriterTo) io.ReadCloser {
return &writerToReaderAdapter{WriterTo: wt}
}

type writerToReaderAdapter struct {
io.WriterTo
prOnce sync.Once
pr *io.PipeReader
}

func (a *writerToReaderAdapter) initPipe() {
r, w := io.Pipe()
go func() {
_, err := a.WriteTo(w)
w.CloseWithError(err)
}()
a.pr = r
}

func (a *writerToReaderAdapter) Read(b []byte) (n int, err error) {
a.prOnce.Do(a.initPipe)
return a.pr.Read(b)
}

func (a *writerToReaderAdapter) Close() error {
a.prOnce.Do(a.initPipe)
return a.pr.Close()
}

type writeCounter struct {
io.Writer
n atomic.Int64
}

func (wc *writeCounter) Count() int64 { return wc.n.Load() }
func (wc *writeCounter) Write(b []byte) (n int, err error) {
n, err = wc.Writer.Write(b)
wc.n.Add(int64(n))
return
}

// WriterToFunc impelments [io.WriterTo]. The io.Writer passed to the function will be wrapped.
type WriterToFunc func(w io.Writer) (err error)

// WriteTo writes to w.
//
// The value in w is wrapped when passed to fn keeping track of how bytes are written by fn.
func (fn WriterToFunc) WriteTo(w io.Writer) (n int64, err error) {
wc := writeCounter{Writer: w}
err = fn(&wc)
return wc.Count(), err
}
49 changes: 49 additions & 0 deletions pkg/skip_response_writer_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
package goa

import (
"bytes"
"io"
"strings"
"testing"
)

func TestSkipResponseWriter(t *testing.T) {
const input = "Hello, World!"
var responseWriter io.ReadCloser

responseWriter = SkipResponseWriter(strings.NewReader(input))
defer func() {
err := responseWriter.Close()
if err != nil {
t.Error(err)
}
}()
_, ok := responseWriter.(io.WriterTo)
if !ok {
t.Errorf("SkipResponseWriter's result must implement io.WriterTo")
}

var writerToBuffer bytes.Buffer
_, err := io.Copy(&writerToBuffer, responseWriter) // io.Copy uses WriterTo if implemented
if err != nil {
t.Fatal(err)
}
if writerToBuffer.String() != input {
t.Errorf("WriteTo: expected=%q actual=%q", input, writerToBuffer.String())
}

responseWriter = SkipResponseWriter(strings.NewReader(input))
defer func() {
err := responseWriter.Close()
if err != nil {
t.Error(err)
}
}()
readBytes, err := io.ReadAll(responseWriter) // io.ReadAll ignores WriterTo and calls Read
if err != nil {
t.Fatal(err)
}
if string(readBytes) != input {
t.Errorf("Read: expected=%q actual=%q", input, string(readBytes))
}
}