optimize response

src_response.go

@@ -8,58 +8,63 @@ import (
 )
 
 type response struct {
-	resp http.Response
+	r http.Response
 }
 
 func newResponse() *response {
-	return &response{
-		resp: http.Response{
-			ProtoMajor: 1,
-			ProtoMinor: 1,
-			Close:      true,
-			Header: http.Header{
-				"Date": []string{time.Now().UTC().Format(http.TimeFormat)},
-			},
+	return &response{http.Response{
+		ProtoMajor: 1,
+		ProtoMinor: 1,
+		Close:      true,
+		Header: http.Header{
+			"Date": []string{time.Now().UTC().Format(http.TimeFormat)},
 		},
-	}
+	}}
 }
 
 func (r *response) Header() http.Header {
-	if r.resp.StatusCode != 0 {
+	if r.r.StatusCode != 0 {
 		// wrote header
 		return http.Header{}
 	}
-	return r.resp.Header
+	return r.r.Header
+}
+
+func (r *response) WriteHeader(statusCode int) {
+	if r.r.StatusCode == 0 {
+		// not wrote header
+		r.r.StatusCode = statusCode
+	}
+}
+
+type respBuf struct {
+	*bytes.Buffer
+	io.Closer
 }
 
 func (r *response) Write(b []byte) (int, error) {
-	if r.resp.Body == nil {
-		r.resp.Body = &respBuf{bytes.NewBuffer(b)}
+	if r.r.Body == nil {
+		r.r.Body = &respBuf{Buffer: bytes.NewBuffer(b)}
 		return len(b), nil
 	}
-	return r.resp.Body.(*respBuf).Write(b)
+	return r.r.Body.(*respBuf).Write(b)
 }
 
 func (r *response) WriteString(s string) (int, error) {
-	if r.resp.Body == nil {
-		r.resp.Body = &respBuf{bytes.NewBufferString(s)}
+	if r.r.Body == nil {
+		r.r.Body = &respBuf{Buffer: bytes.NewBufferString(s)}
 		return len(s), nil
 	}
-	return r.resp.Body.(*respBuf).WriteString(s)
-}
-
-func (r *response) WriteHeader(statusCode int) {
-	if r.resp.StatusCode == 0 {
-		// not wrote header
-		r.resp.StatusCode = statusCode
-	}
+	return r.r.Body.(*respBuf).WriteString(s)
 }
 
 // flush flushes buffered data to the client.
-func (r *response) flush(w io.Writer) error {
+func (r *response) flush(w io.WriteCloser) error {
 	r.WriteHeader(400)
-	if r.resp.Body != nil {
-		r.resp.ContentLength = int64(r.resp.Body.(*respBuf).Len())
+	if r.r.Body != nil {
+		b := r.r.Body.(*respBuf)
+		r.r.ContentLength = int64(b.Len())
+		b.Closer = w
 	}
-	return r.resp.Write(w)
+	return r.r.Write(w)
 }