1. docker
  2. Commit 24310b5b4ab600ab5cb046d2f2ceaaa086b30be3
Browse code

Add BytesPipe datastructure to ioutils

Signed-off-by: Alexander Morozov <lk4d4@docker.com>

Alexander Morozov authored on 2015/09/10 01:47:24
Showing 2 changed files
pkg/ioutils/bytespipe.go
History View file @ 24310b5
1 1 
new file mode 100644
... ... 
@@ -0,0 +1,82 @@
0 
+package ioutils
1 
+
2 
+const maxCap = 10 * 1e6
3 
+
4 
+// BytesPipe is io.ReadWriter which works similary to pipe(queue).
5 
+// All written data could be read only once. Also BytesPipe trying to adjust
6 
+// internal []byte slice to current needs, so there won't be overgrown buffer
7 
+// after highload peak.
8 
+// BytesPipe isn't goroutine-safe, caller must synchronize it if needed.
9 
+type BytesPipe struct {
10 
+	buf      []byte
11 
+	lastRead int
12 
+}
13 
+
14 
+// NewBytesPipe creates new BytesPipe, initialized by specified slice.
15 
+// If buf is nil, then it will be initialized with slice which cap is 64.
16 
+// buf will be adjusted in a way that len(buf) == 0, cap(buf) == cap(buf).
17 
+func NewBytesPipe(buf []byte) *BytesPipe {
18 
+	if cap(buf) == 0 {
19 
+		buf = make([]byte, 0, 64)
20 
+	}
21 
+	return &BytesPipe{
22 
+		buf: buf[:0],
23 
+	}
24 
+}
25 
+
26 
+func (bp *BytesPipe) grow(n int) {
27 
+	if len(bp.buf)+n > cap(bp.buf) {
28 
+		// not enough space
29 
+		var buf []byte
30 
+		remain := bp.len()
31 
+		if remain+n <= cap(bp.buf)/2 {
32 
+			// enough space in current buffer, just move data to head
33 
+			copy(bp.buf, bp.buf[bp.lastRead:])
34 
+			buf = bp.buf[:remain]
35 
+		} else {
36 
+			// reallocate buffer
37 
+			buf = make([]byte, remain, 2*cap(bp.buf)+n)
38 
+			copy(buf, bp.buf[bp.lastRead:])
39 
+		}
40 
+		bp.buf = buf
41 
+		bp.lastRead = 0
42 
+	}
43 
+}
44 
+
45 
+// Write writes p to BytesPipe.
46 
+// It can increase cap of internal []byte slice in a process of writing.
47 
+func (bp *BytesPipe) Write(p []byte) (n int, err error) {
48 
+	bp.grow(len(p))
49 
+	bp.buf = append(bp.buf, p...)
50 
+	return
51 
+}
52 
+
53 
+func (bp *BytesPipe) len() int {
54 
+	return len(bp.buf) - bp.lastRead
55 
+}
56 
+
57 
+func (bp *BytesPipe) crop() {
58 
+	// shortcut for empty buffer
59 
+	if bp.lastRead == len(bp.buf) {
60 
+		bp.lastRead = 0
61 
+		bp.buf = bp.buf[:0]
62 
+	}
63 
+	r := bp.len()
64 
+	// if we have too large buffer for too small data
65 
+	if cap(bp.buf) > maxCap && r < cap(bp.buf)/10 {
66 
+		copy(bp.buf, bp.buf[bp.lastRead:])
67 
+		// will use same underlying slice until reach cap
68 
+		bp.buf = bp.buf[:r : cap(bp.buf)/2]
69 
+		bp.lastRead = 0
70 
+	}
71 
+}
72 
+
73 
+// Read reads bytes from BytesPipe.
74 
+// Data could be read only once.
75 
+// Internal []byte slice could be shrinked.
76 
+func (bp *BytesPipe) Read(p []byte) (n int, err error) {
77 
+	n = copy(p, bp.buf[bp.lastRead:])
78 
+	bp.lastRead += n
79 
+	bp.crop()
80 
+	return
81 
+}
pkg/ioutils/bytespipe_test.go
History View file @ 24310b5
0 82 
new file mode 100644
... ... 
@@ -0,0 +1,81 @@
0 
+package ioutils
1 
+
2 
+import "testing"
3 
+
4 
+func TestBytesPipeRead(t *testing.T) {
5 
+	buf := NewBytesPipe(nil)
6 
+	buf.Write([]byte("12"))
7 
+	buf.Write([]byte("34"))
8 
+	buf.Write([]byte("56"))
9 
+	buf.Write([]byte("78"))
10 
+	buf.Write([]byte("90"))
11 
+	rd := make([]byte, 4)
12 
+	n, err := buf.Read(rd)
13 
+	if err != nil {
14 
+		t.Fatal(err)
15 
+	}
16 
+	if n != 4 {
17 
+		t.Fatalf("Wrong number of bytes read: %d, should be %d", n, 4)
18 
+	}
19 
+	if string(rd) != "1234" {
20 
+		t.Fatalf("Read %s, but must be %s", rd, "1234")
21 
+	}
22 
+	n, err = buf.Read(rd)
23 
+	if err != nil {
24 
+		t.Fatal(err)
25 
+	}
26 
+	if n != 4 {
27 
+		t.Fatalf("Wrong number of bytes read: %d, should be %d", n, 4)
28 
+	}
29 
+	if string(rd) != "5678" {
30 
+		t.Fatalf("Read %s, but must be %s", rd, "5679")
31 
+	}
32 
+	n, err = buf.Read(rd)
33 
+	if err != nil {
34 
+		t.Fatal(err)
35 
+	}
36 
+	if n != 2 {
37 
+		t.Fatalf("Wrong number of bytes read: %d, should be %d", n, 2)
38 
+	}
39 
+	if string(rd[:n]) != "90" {
40 
+		t.Fatalf("Read %s, but must be %s", rd, "90")
41 
+	}
42 
+}
43 
+
44 
+func TestBytesPipeWrite(t *testing.T) {
45 
+	buf := NewBytesPipe(nil)
46 
+	buf.Write([]byte("12"))
47 
+	buf.Write([]byte("34"))
48 
+	buf.Write([]byte("56"))
49 
+	buf.Write([]byte("78"))
50 
+	buf.Write([]byte("90"))
51 
+	if string(buf.buf) != "1234567890" {
52 
+		t.Fatalf("Buffer %s, must be %s", buf.buf, "1234567890")
53 
+	}
54 
+}
55 
+
56 
+func BenchmarkBytesPipeWrite(b *testing.B) {
57 
+	for i := 0; i < b.N; i++ {
58 
+		buf := NewBytesPipe(nil)
59 
+		for j := 0; j < 1000; j++ {
60 
+			buf.Write([]byte("pretty short line, because why not?"))
61 
+		}
62 
+	}
63 
+}
64 
+
65 
+func BenchmarkBytesPipeRead(b *testing.B) {
66 
+	rd := make([]byte, 1024)
67 
+	for i := 0; i < b.N; i++ {
68 
+		b.StopTimer()
69 
+		buf := NewBytesPipe(nil)
70 
+		for j := 0; j < 1000; j++ {
71 
+			buf.Write(make([]byte, 1024))
72 
+		}
73 
+		b.StartTimer()
74 
+		for j := 0; j < 1000; j++ {
75 
+			if n, _ := buf.Read(rd); n != 1024 {
76 
+				b.Fatalf("Wrong number of bytes: %d", n)
77 
+			}
78 
+		}
79 
+	}
80 
+}