Merge pull request #297 from marco6/improve-get-blob

Improve Message.getBlob performance

internal/protocol/message.go

@@ -322,18 +322,9 @@ func (m *Message) getString() string {
 func (m *Message) getBlob() []byte {
 	size := m.getUint64()
 	data := make([]byte, size)
-	for i := range data {
-		data[i] = m.getUint8()
-	}
-	pad := 0
-	if (size % messageWordSize) != 0 {
-		// Account for padding
-		pad = int(messageWordSize - (size % messageWordSize))
-	}
-	// Consume padding
-	for i := 0; i < pad; i++ {
-		m.getUint8()
-	}
+	b := m.bufferForGet()
+	defer b.Advance(int(alignUp(size, messageWordSize)))
+	copy(data, b.Bytes[b.Offset:])
 	return data
 }
 
@@ -684,3 +675,8 @@ func (r *Rows) ColumnTypes() ([]string, error) {
 
 	return kinds, err
 }
+
+// alignUp rounds n up to a multiple of a. a must be a power of 2.
+func alignUp(n, a uint64) uint64 {
+	return (n + a - 1) &^ (a - 1)
+}

internal/protocol/message_internal_test.go

@@ -302,6 +302,31 @@ func TestMessage_getBlob(t *testing.T) {
 	}
 }
 
+func BenchmarkMessage_getBlob(b *testing.B) {
+	makeBlob := func(size int) []byte {
+		blob := make([]byte, size)
+		for i := range blob {
+			blob[i] = byte(i)
+		}
+		return blob
+	}
+
+	for _, size := range []int{16, 64, 256, 1024, 4096, 8096} {
+		b.Run(fmt.Sprintf("%d", size), func(b *testing.B) {
+			message := Message{}
+			message.Init(size + 16)
+
+			message.putBlob(makeBlob(size))
+			message.putHeader(0, 0)
+
+			for i := 0; i < b.N; i++ {
+				message.Rewind()
+				_ = message.getBlob()
+			}
+		})
+	}
+}
+
 // The overflowing string ends exactly at word boundary.
 func TestMessage_getString_Overflow_WordBoundary(t *testing.T) {
 	message := Message{}