From c1d3705be053527784481ff49a2f55ed86bbfa3d Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Sat, 28 Sep 2019 20:38:24 +0300
Subject: [PATCH] app/vmselect/netstorage: marshal block outside
 tmpBlocksFile.WriteBlock

This allows re-using the destination buffer for marshaling in the outer loop.
---
 app/vmselect/netstorage/netstorage.go           |  5 ++++-
 app/vmselect/netstorage/tmp_blocks_file.go      | 16 ++++++----------
 app/vmselect/netstorage/tmp_blocks_file_test.go |  5 ++++-
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/app/vmselect/netstorage/netstorage.go b/app/vmselect/netstorage/netstorage.go
index 4e223b6bb..f52d5c3dd 100644
--- a/app/vmselect/netstorage/netstorage.go
+++ b/app/vmselect/netstorage/netstorage.go
@@ -484,9 +484,12 @@ func ProcessSearchQuery(sq *storage.SearchQuery, fetchData bool, deadline Deadli
 	tbf := getTmpBlocksFile()
 	m := make(map[string][]tmpBlockAddr)
 	blocksRead := 0
+	bb := tmpBufPool.Get()
+	defer tmpBufPool.Put(bb)
 	for sr.NextMetricBlock() {
 		blocksRead++
-		addr, err := tbf.WriteBlock(sr.MetricBlock.Block)
+		bb.B = storage.MarshalBlock(bb.B[:0], sr.MetricBlock.Block)
+		addr, err := tbf.WriteBlockData(bb.B)
 		if err != nil {
 			putTmpBlocksFile(tbf)
 			return nil, fmt.Errorf("cannot write data block #%d to temporary blocks file: %s", blocksRead, err)
diff --git a/app/vmselect/netstorage/tmp_blocks_file.go b/app/vmselect/netstorage/tmp_blocks_file.go
index ccf68259e..4e23f03fc 100644
--- a/app/vmselect/netstorage/tmp_blocks_file.go
+++ b/app/vmselect/netstorage/tmp_blocks_file.go
@@ -82,22 +82,18 @@ func (addr tmpBlockAddr) String() string {
 
 var tmpBlocksFilesCreated = metrics.NewCounter(`vm_tmp_blocks_files_created_total`)
 
-// WriteBlock writes b to tbf.
+// WriteBlockData writes b to tbf.
 //
 // It returns errors since the operation may fail on space shortage
 // and this must be handled.
-func (tbf *tmpBlocksFile) WriteBlock(b *storage.Block) (tmpBlockAddr, error) {
-	bb := tmpBufPool.Get()
-	defer tmpBufPool.Put(bb)
-	bb.B = storage.MarshalBlock(bb.B[:0], b)
-
+func (tbf *tmpBlocksFile) WriteBlockData(b []byte) (tmpBlockAddr, error) {
 	var addr tmpBlockAddr
 	addr.offset = tbf.offset
-	addr.size = len(bb.B)
+	addr.size = len(b)
 	tbf.offset += uint64(addr.size)
-	if len(tbf.buf)+len(bb.B) <= cap(tbf.buf) {
+	if len(tbf.buf)+len(b) <= cap(tbf.buf) {
 		// Fast path - the data fits tbf.buf
-		tbf.buf = append(tbf.buf, bb.B...)
+		tbf.buf = append(tbf.buf, b...)
 		return addr, nil
 	}
 
@@ -111,7 +107,7 @@ func (tbf *tmpBlocksFile) WriteBlock(b *storage.Block) (tmpBlockAddr, error) {
 		tmpBlocksFilesCreated.Inc()
 	}
 	_, err := tbf.f.Write(tbf.buf)
-	tbf.buf = append(tbf.buf[:0], bb.B...)
+	tbf.buf = append(tbf.buf[:0], b...)
 	if err != nil {
 		return addr, fmt.Errorf("cannot write block to %q: %s", tbf.f.Name(), err)
 	}
diff --git a/app/vmselect/netstorage/tmp_blocks_file_test.go b/app/vmselect/netstorage/tmp_blocks_file_test.go
index 586ce80f8..eba2292a9 100644
--- a/app/vmselect/netstorage/tmp_blocks_file_test.go
+++ b/app/vmselect/netstorage/tmp_blocks_file_test.go
@@ -77,9 +77,12 @@ func testTmpBlocksFile() error {
 			// Write blocks until their summary size exceeds `size`.
 			var addrs []tmpBlockAddr
 			var blocks []*storage.Block
+			bb := tmpBufPool.Get()
+			defer tmpBufPool.Put(bb)
 			for tbf.offset < uint64(size) {
 				b := createBlock()
-				addr, err := tbf.WriteBlock(b)
+				bb.B = storage.MarshalBlock(bb.B[:0], b)
+				addr, err := tbf.WriteBlockData(bb.B)
 				if err != nil {
 					return fmt.Errorf("cannot write block at offset %d: %s", tbf.offset, err)
 				}