From 33d531dfa49b7477dd51b54b870508b8b7eafee2 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Mon, 19 Aug 2013 22:56:54 -0700
Subject: [PATCH] net/http: support WriteString on the ResponseWriter

Fixes #5377

R=golang-dev, r
CC=golang-dev
https://golang.org/cl/12991046
---
 src/pkg/net/http/serve_test.go | 15 +++++++++++++++
 src/pkg/net/http/server.go     | 19 ++++++++++++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/pkg/net/http/serve_test.go b/src/pkg/net/http/serve_test.go
index df4367e2b2..5d08d2faca 100644
--- a/src/pkg/net/http/serve_test.go
+++ b/src/pkg/net/http/serve_test.go
@@ -1994,6 +1994,21 @@ func TestNoContentTypeOnNotModified(t *testing.T) {
 	}
 }
 
+func TestResponseWriterWriteStringAllocs(t *testing.T) {
+	ht := newHandlerTest(HandlerFunc(func(w ResponseWriter, r *Request) {
+		if r.URL.Path == "/s" {
+			io.WriteString(w, "Hello world")
+		} else {
+			w.Write([]byte("Hello world"))
+		}
+	}))
+	before := testing.AllocsPerRun(25, func() { ht.rawResponse("GET / HTTP/1.0") })
+	after := testing.AllocsPerRun(25, func() { ht.rawResponse("GET /s HTTP/1.0") })
+	if int(after) >= int(before) {
+		t.Errorf("WriteString allocs of %v >= Write allocs of %v", after, before)
+	}
+}
+
 func BenchmarkClientServer(b *testing.B) {
 	b.ReportAllocs()
 	b.StopTimer()
diff --git a/src/pkg/net/http/server.go b/src/pkg/net/http/server.go
index 9702aee274..3c327839cb 100644
--- a/src/pkg/net/http/server.go
+++ b/src/pkg/net/http/server.go
@@ -953,6 +953,15 @@ func (w *response) bodyAllowed() bool {
 // bufferBeforeChunkingSize smaller and having bufio's fast-paths deal
 // with this instead.
 func (w *response) Write(data []byte) (n int, err error) {
+	return w.write(len(data), data, "")
+}
+
+func (w *response) WriteString(data string) (n int, err error) {
+	return w.write(len(data), nil, data)
+}
+
+// either dataB or dataS is non-zero.
+func (w *response) write(lenData int, dataB []byte, dataS string) (n int, err error) {
 	if w.conn.hijacked() {
 		log.Print("http: response.Write on hijacked connection")
 		return 0, ErrHijacked
@@ -960,18 +969,22 @@ func (w *response) Write(data []byte) (n int, err error) {
 	if !w.wroteHeader {
 		w.WriteHeader(StatusOK)
 	}
-	if len(data) == 0 {
+	if lenData == 0 {
 		return 0, nil
 	}
 	if !w.bodyAllowed() {
 		return 0, ErrBodyNotAllowed
 	}
 
-	w.written += int64(len(data)) // ignoring errors, for errorKludge
+	w.written += int64(lenData) // ignoring errors, for errorKludge
 	if w.contentLength != -1 && w.written > w.contentLength {
 		return 0, ErrContentLength
 	}
-	return w.w.Write(data)
+	if dataB != nil {
+		return w.w.Write(dataB)
+	} else {
+		return w.w.WriteString(dataS)
+	}
 }
 
 func (w *response) finishRequest() {