Exported ByteBuffer

bytebuffer.go

@@ -0,0 +1,46 @@
+package fasthttp
+
+import (
+	"sync"
+)
+
+const (
+	defaultByteBufferSize = 128
+)
+
+// ByteBuffer provides byte buffer, which can be used with fasthttp API
+// in order to minimize memory allocations.
+//
+// ByteBuffer may be used with functions appending data to the given []byte
+// slice. See example code for details.
+//
+// Use AcquireByteBuffer for obtaining an empty byte buffer.
+type ByteBuffer struct {
+	B []byte
+}
+
+// AcquireByteBuffer returns an empty byte buffer from the pool.
+//
+// Acquired byte buffer may be returned to the pool via ReleaseByteBuffer call.
+// This reduces the number of memory allocations required for byte buffer
+// management.
+func AcquireByteBuffer() *ByteBuffer {
+	v := byteBufferPool.Get()
+	if v == nil {
+		return &ByteBuffer{
+			B: make([]byte, 0, defaultByteBufferSize),
+		}
+	}
+	return v.(*ByteBuffer)
+}
+
+// ReleaseByteBuffer returns byte buffer to the pool.
+//
+// ByteBuffer.B mustn't be touched after returning it to the pool.
+// Otherwise data races occur.
+func ReleaseByteBuffer(b *ByteBuffer) {
+	b.B = b.B[:0]
+	byteBufferPool.Put(b)
+}
+
+var byteBufferPool sync.Pool

bytebuffer_example_test.go

@@ -0,0 +1,29 @@
+package fasthttp_test
+
+import (
+	"fmt"
+
+	"github.com/valyala/fasthttp"
+)
+
+func ExampleByteBuffer() {
+	// This request handler sets 'Your-IP' response header
+	// to 'Your IP is <ip>'. It uses ByteBuffer for constructing response
+	// header value with zero memory allocations.
+	yourIPRequestHandler := func(ctx *fasthttp.RequestCtx) {
+		b := fasthttp.AcquireByteBuffer()
+		b.B = append(b.B, "Your IP is <"...)
+		b.B = fasthttp.AppendIPv4(b.B, ctx.RemoteIP())
+		b.B = append(b.B, ">"...)
+		ctx.Response.Header.SetBytesV("Your-IP", b.B)
+
+		fmt.Fprintf(ctx, "Check response headers - they must contain 'Your-IP: %s", b.B)
+
+		// It is safe to release byte buffer now, since it is
+		// no longer used.
+		fasthttp.ReleaseByteBuffer(b)
+	}
+
+	// Start fasthttp server returning your ip in response headers.
+	fasthttp.ListenAndServe(":8080", yourIPRequestHandler)
+}

bytebuffer_test.go

@@ -0,0 +1,43 @@
+package fasthttp
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestByteBufferAcquireReleaseSerial(t *testing.T) {
+	testByteBufferAcquireRelease(t)
+}
+
+func TestByteBufferAcquireReleaseConcurrent(t *testing.T) {
+	concurrency := 10
+	ch := make(chan struct{}, concurrency)
+	for i := 0; i < concurrency; i++ {
+		go func() {
+			testByteBufferAcquireRelease(t)
+			ch <- struct{}{}
+		}()
+	}
+
+	for i := 0; i < concurrency; i++ {
+		select {
+		case <-ch:
+		case <-time.After(time.Second):
+			t.Fatalf("timeout!")
+		}
+	}
+}
+
+func testByteBufferAcquireRelease(t *testing.T) {
+	for i := 0; i < 10; i++ {
+		b := AcquireByteBuffer()
+		b.B = append(b.B, "num "...)
+		b.B = AppendUint(b.B, i)
+		expectedS := fmt.Sprintf("num %d", i)
+		if string(b.B) != expectedS {
+			t.Fatalf("unexpected result: %q. Expecting %q", b.B, expectedS)
+		}
+		ReleaseByteBuffer(b)
+	}
+}

fs.go

@@ -106,12 +106,12 @@ func NewVHostPathRewriter(slashesCount int) PathRewriteFunc {
 		if len(host) == 0 {
 			host = strInvalidHost
 		}
-		b := acquireByteBuffer()
-		b.b = append(b.b, '/')
-		b.b = append(b.b, host...)
-		b.b = append(b.b, path...)
-		ctx.URI().SetPathBytes(b.b)
-		releaseByteBuffer(b)
+		b := AcquireByteBuffer()
+		b.B = append(b.B, '/')
+		b.B = append(b.B, host...)
+		b.B = append(b.B, path...)
+		ctx.URI().SetPathBytes(b.B)
+		ReleaseByteBuffer(b)
 
 		return ctx.Path()
 	}
@@ -119,27 +119,6 @@ func NewVHostPathRewriter(slashesCount int) PathRewriteFunc {
 
 var strInvalidHost = []byte("invalid-host")
 
-func acquireByteBuffer() *byteBuffer {
-	return byteBufferPool.Get().(*byteBuffer)
-}
-
-func releaseByteBuffer(b *byteBuffer) {
-	b.b = b.b[:0]
-	byteBufferPool.Put(b)
-}
-
-var byteBufferPool = &sync.Pool{
-	New: func() interface{} {
-		return &byteBuffer{
-			b: make([]byte, 0, 128),
-		}
-	},
-}
-
-type byteBuffer struct {
-	b []byte
-}
-
 // NewPathSlashesStripper returns path rewriter, which strips slashesCount
 // leading slashes from the path.
 //