Repository: goadesign/goa
Pull request #3537 (2024-06-18 23:44:31)

SkipResponseWriter
Implementation for #3522

Since the tests, don't include the comment, here's the result:

```go
// Service is the admin service interface.
type Service interface {
	// Assets implements assets.

	// If body implements [io.WriterTo], that implementation will be used instead.
	// Consider [goa.design/goa/v3/pkg.SkipResponseWriter] to adapt existing
	// implementations.
	Assets(context.Context, *AuthenticatedRequest) (res *CustomJSONResponse, body io.ReadCloser, err error)
}
```

Note: Tried this out in one of my projects and added `WriterToFunc` to the exported API as well, since I found myself wanting it for most of the handlers I was rewriting. Allows you to use a closure and avoids needing to manually keep track of bytes written, while still preserving the `io.WriterTo` interface. Useful for times when you're passing the writer to something else that doesn't report how many bytes are written like `json.Encoder` or `xml.Encoder`.

raphael: Thank you for the great contribution! I was wondering if you'd be up to updating the [upload_download](https://github.com/goadesign/examples/tree/master/upload_download) example to take advantage of this new feature after we release it?

Patch:
diff --git a/codegen/service/templates/service.go.tpl b/codegen/service/templates/service.go.tpl
index a12efc67..5164c620 100644
--- a/codegen/service/templates/service.go.tpl
+++ b/codegen/service/templates/service.go.tpl
@@ -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" }}
diff --git a/http/codegen/templates/server_handler_init.go.tpl b/http/codegen/templates/server_handler_init.go.tpl
index 254857dc..00fb7094 100644
--- a/http/codegen/templates/server_handler_init.go.tpl
+++ b/http/codegen/templates/server_handler_init.go.tpl
@@ -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
+				}
+			}
+			return
+		}
 		// handle immediate read error like a returned error
 		buf := bufio.NewReader(o.Body)
 		if _, err := buf.Peek(1); err != nil && err != io.EOF {
diff --git a/pkg/skip_response_writer.go b/pkg/skip_response_writer.go
new file mode 100644
index 00000000..1800257c
--- /dev/null
+++ b/pkg/skip_response_writer.go
@@ -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
+}


Test patch:
diff --git a/http/codegen/testdata/handler_init_functions.go b/http/codegen/testdata/handler_init_functions.go
index 657d5bf6..5450611a 100644
--- a/http/codegen/testdata/handler_init_functions.go
+++ b/http/codegen/testdata/handler_init_functions.go
@@ -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 {
diff --git a/pkg/skip_response_writer_test.go b/pkg/skip_response_writer_test.go
new file mode 100644
index 00000000..75dde0bc
--- /dev/null
+++ b/pkg/skip_response_writer_test.go
@@ -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))
+	}
+}