feat(rpc): response implements interface read from

close: cubefs#2652

Signed-off-by: slasher <[email protected]>

blobstore/common/rpc/auditlog/auditlog_test.go

@@ -15,9 +15,11 @@
 package auditlog
 
 import (
+	"bytes"
 	"encoding/json"
 	"encoding/xml"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"math/rand"
 	"net/http"
@@ -74,10 +76,34 @@ func initServer(t *testing.T, name string, cfg Config) (server *httptest.Server,
 		extraHeader.Set("Extra-header2", "header2 value")
 		extraHeader.Add("Extra-header2", "header2 value2")
 
-		w.WriteHeader(http.StatusOK)
 		data, err := json.Marshal(testRespData{Result: "success"})
 		require.NoError(t, err)
 		w.Write(data)
+
+		rw := w.(*responseWriter)
+		if !rw.no2xxBody && rw.bodyLimit > 0 {
+			require.Equal(t, len(data), rw.n)
+			require.Equal(t, data, rw.body[:rw.n])
+		} else {
+			require.Equal(t, 0, rw.n)
+		}
+	}
+
+	streamHandler := func(w http.ResponseWriter, req *http.Request) {
+		size := int64(64 * 1024)
+		buffer := make([]byte, size)
+		for range [1024]struct{}{} {
+			_, err := io.CopyN(w, bytes.NewReader(buffer), size)
+			require.NoError(t, err)
+		}
+
+		span := trace.SpanFromContextSafe(req.Context())
+		span.SetTag("response", "readfrom")
+		span.AppendRPCTrackLog([]string{"stream"})
+
+		rw := w.(*responseWriter)
+		require.Equal(t, 0, rw.n)
+		require.Equal(t, 1024*size, rw.bodyWritten)
 	}
 
 	errorResponseHandler := func(w http.ResponseWriter, req *http.Request) {
@@ -97,6 +123,8 @@ func initServer(t *testing.T, name string, cfg Config) (server *httptest.Server,
 		switch req.URL.Path {
 		case "/":
 			ah.Handler(w, req, bussinessHandler)
+		case "/stream":
+			ah.Handler(w, req, streamHandler)
 		case "/error-response":
 			ah.Handler(w, req, errorResponseHandler)
 		default:
@@ -277,6 +305,24 @@ func TestNoLogBody(t *testing.T) {
 	require.NoError(t, open.Close())
 }
 
+func TestResponseReadFrom(t *testing.T) {
+	server, tmpDir, lc := initServer(t, "testResponseReadFrom", Config{LogFormat: LogFormatJSON})
+	defer func() {
+		server.Close()
+		os.RemoveAll(tmpDir)
+		lc.Close()
+	}()
+
+	req, err := http.NewRequest(http.MethodGet, server.URL+"/stream", nil)
+	require.NoError(t, err)
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	_, err = io.Copy(io.Discard, resp.Body)
+	require.NoError(t, err)
+}
+
 func TestBodylimit(t *testing.T) {
 	for _, limit := range []struct {
 		actual, expected int

blobstore/common/rpc/auditlog/response.go

@@ -16,6 +16,7 @@ package auditlog
 
 import (
 	"bufio"
+	"io"
 	"net"
 	"net/http"
 	"strconv"
@@ -59,6 +60,23 @@ func (w *responseWriter) Write(b []byte) (int, error) {
 	return n, err
 }
 
+// ReadFrom implement io.ReaderFrom when io.Copy.
+// Response with this function will not hold first body bytes in local buffer.
+func (w *responseWriter) ReadFrom(src io.Reader) (n int64, err error) {
+	if !w.hasWroteHeader {
+		w.WriteHeader(http.StatusOK)
+		w.hasWroteHeader = true
+	}
+	if rf, ok := w.ResponseWriter.(io.ReaderFrom); ok {
+		n, err = rf.ReadFrom(src)
+		w.bodyWritten += int64(n)
+		return
+	}
+	n, err = io.Copy(w.ResponseWriter, src)
+	w.bodyWritten += int64(n)
+	return
+}
+
 func (w *responseWriter) WriteHeader(code int) {
 	if w.hasWroteHeader {
 		return