From 3c79a0a23c1940ab4479332fb3e0127265650ce3 Mon Sep 17 00:00:00 2001
From: Ewan Chou <coocood@gmail.com>
Date: Tue, 26 Feb 2019 13:05:06 +0800
Subject: [PATCH] add GetWithBuf method to avoid memory allocation

---
 cache.go      | 15 +++++++++++++--
 cache_test.go | 21 ++++++++++++++++++++-
 segment.go    |  8 ++++++--
 3 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/cache.go b/cache.go
index 04049af..9a031c0 100644
--- a/cache.go
+++ b/cache.go
@@ -60,7 +60,18 @@ func (cache *Cache) Get(key []byte) (value []byte, err error) {
 	hashVal := hashFunc(key)
 	segID := hashVal & segmentAndOpVal
 	cache.locks[segID].Lock()
-	value, _, err = cache.segments[segID].get(key, hashVal)
+	value, _, err = cache.segments[segID].get(key, nil, hashVal)
+	cache.locks[segID].Unlock()
+	return
+}
+
+// GetWithBuf copies the value to the buf or returns not found error.
+// This method doesn't allocate memory when the capacity of buf is greater or equal to value.
+func (cache *Cache) GetWithBuf(key, buf []byte) (value []byte, err error) {
+	hashVal := hashFunc(key)
+	segID := hashVal & segmentAndOpVal
+	cache.locks[segID].Lock()
+	value, _, err = cache.segments[segID].get(key, buf, hashVal)
 	cache.locks[segID].Unlock()
 	return
 }
@@ -70,7 +81,7 @@ func (cache *Cache) GetWithExpiration(key []byte) (value []byte, expireAt uint32
 	hashVal := hashFunc(key)
 	segID := hashVal & segmentAndOpVal
 	cache.locks[segID].Lock()
-	value, expireAt, err = cache.segments[segID].get(key, hashVal)
+	value, expireAt, err = cache.segments[segID].get(key, nil, hashVal)
 	cache.locks[segID].Unlock()
 	return
 }
diff --git a/cache_test.go b/cache_test.go
index b00b217..763929a 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -511,12 +511,14 @@ func BenchmarkMapSet(b *testing.B) {
 }
 
 func BenchmarkCacheGet(b *testing.B) {
+	b.ReportAllocs()
 	b.StopTimer()
 	cache := NewCache(256 * 1024 * 1024)
 	var key [8]byte
+	buf := make([]byte, 64)
 	for i := 0; i < b.N; i++ {
 		binary.LittleEndian.PutUint64(key[:], uint64(i))
-		cache.Set(key[:], make([]byte, 8), 0)
+		cache.Set(key[:], buf, 0)
 	}
 	b.StartTimer()
 	for i := 0; i < b.N; i++ {
@@ -525,6 +527,23 @@ func BenchmarkCacheGet(b *testing.B) {
 	}
 }
 
+func BenchmarkCacheGetWithBuf(b *testing.B) {
+	b.ReportAllocs()
+	b.StopTimer()
+	cache := NewCache(256 * 1024 * 1024)
+	var key [8]byte
+	buf := make([]byte, 64)
+	for i := 0; i < b.N; i++ {
+		binary.LittleEndian.PutUint64(key[:], uint64(i))
+		cache.Set(key[:], buf, 0)
+	}
+	b.StartTimer()
+	for i := 0; i < b.N; i++ {
+		binary.LittleEndian.PutUint64(key[:], uint64(i))
+		cache.GetWithBuf(key[:], buf)
+	}
+}
+
 func BenchmarkCacheGetWithExpiration(b *testing.B) {
 	b.StopTimer()
 	cache := NewCache(256 * 1024 * 1024)
diff --git a/segment.go b/segment.go
index f929954..9568836 100644
--- a/segment.go
+++ b/segment.go
@@ -189,7 +189,7 @@ func (seg *segment) evacuate(entryLen int64, slotId uint8, now uint32) (slotModi
 	return
 }
 
-func (seg *segment) get(key []byte, hashVal uint64) (value []byte, expireAt uint32, err error) {
+func (seg *segment) get(key, buf []byte, hashVal uint64) (value []byte, expireAt uint32, err error) {
 	slotId := uint8(hashVal >> 8)
 	hash16 := uint16(hashVal >> 16)
 	slot := seg.getSlot(slotId)
@@ -217,7 +217,11 @@ func (seg *segment) get(key []byte, hashVal uint64) (value []byte, expireAt uint
 	atomic.AddInt64(&seg.totalTime, int64(now-hdr.accessTime))
 	hdr.accessTime = now
 	seg.rb.WriteAt(hdrBuf[:], ptr.offset)
-	value = make([]byte, hdr.valLen)
+	if cap(buf) >= int(hdr.valLen) {
+		value = buf[:hdr.valLen]
+	} else {
+		value = make([]byte, hdr.valLen)
+	}
 
 	seg.rb.ReadAt(value, ptr.offset+ENTRY_HDR_SIZE+int64(hdr.keyLen))
 	atomic.AddInt64(&seg.hitCount, 1)