From a47e26eb67ac2657e4b5a62b1975bb2b65e0b8b3 Mon Sep 17 00:00:00 2001
From: Ewan Chou <coocood@gmail.com>
Date: Sat, 27 May 2017 10:51:13 +0800
Subject: [PATCH] support iterator

---
 README.md     |  6 ++--
 cache_test.go | 28 ++++++++++++++++++
 iterator.go   | 80 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+), 2 deletions(-)
 create mode 100644 iterator.go

diff --git a/README.md b/README.md
index 98ada37..a50760f 100644
--- a/README.md
+++ b/README.md
@@ -17,6 +17,7 @@ without increased latency and degraded throughput.
 * Nearly LRU algorithm
 * Strictly limited memory usage
 * Come with a toy server that supports a few basic Redis commands with pipeline
+* Iterator support
 
 ## Performance
 
@@ -51,8 +52,8 @@ fmt.Println("entry count ", cache.EntryCount())
 
 ## Notice
 
-* Memory is preallocated. 
-* If you allocate large amount of memory, you may need to set `debug.SetGCPercent()` 
+* Memory is preallocated.
+* If you allocate large amount of memory, you may need to set `debug.SetGCPercent()`
 to a much lower percentage to get a normal GC frequency.
 
 ## How it is done
@@ -70,4 +71,5 @@ Each segment has its own lock, so it supports high concurrent access.
 * Support resize cache size at runtime.
 
 ## License
+
 The MIT License
diff --git a/cache_test.go b/cache_test.go
index 272d926..c8a0bd9 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -304,6 +304,34 @@ func TestInt64Key(t *testing.T) {
 	}
 }
 
+func TestIterator(t *testing.T) {
+	cache := NewCache(1024)
+	count := 10000
+	for i := 0; i < count; i++ {
+		err := cache.Set([]byte(fmt.Sprintf("%d", i)), []byte(fmt.Sprintf("val%d", i)), 0)
+		if err != nil {
+			t.Error(err)
+		}
+	}
+	// Set some value that expires to make sure expired entry is not returned.
+	cache.Set([]byte("abc"), []byte("def"), 1)
+	time.Sleep(2 * time.Second)
+	it := cache.NewIterator()
+	for i := 0; i < count; i++ {
+		entry := it.Next()
+		if entry == nil {
+			t.Fatalf("entry is nil for %d", i)
+		}
+		if string(entry.Value) != "val"+string(entry.Key) {
+			t.Fatalf("entry key value not match %s %s", entry.Key, entry.Value)
+		}
+	}
+	e := it.Next()
+	if e != nil {
+		t.Fail()
+	}
+}
+
 func BenchmarkCacheSet(b *testing.B) {
 	cache := NewCache(256 * 1024 * 1024)
 	var key [8]byte
diff --git a/iterator.go b/iterator.go
new file mode 100644
index 0000000..de07800
--- /dev/null
+++ b/iterator.go
@@ -0,0 +1,80 @@
+package freecache
+
+import (
+	"time"
+	"unsafe"
+)
+
+// Iterator iterates the entries for the cache.
+type Iterator struct {
+	cache      *Cache
+	segmentIdx int
+	slotIdx    int
+	entryIdx   int
+}
+
+// Entry represents a key/value pair.
+type Entry struct {
+	Key   []byte
+	Value []byte
+}
+
+// Next returns the next entry for the iterator.
+// The order of the entries is not guaranteed.
+// If there is no more entries to return, nil will be returned.
+func (it *Iterator) Next() *Entry {
+	for it.segmentIdx < 256 {
+		entry := it.nextForSegment(it.segmentIdx)
+		if entry != nil {
+			return entry
+		}
+		it.segmentIdx++
+		it.slotIdx = 0
+		it.entryIdx = 0
+	}
+	return nil
+}
+
+func (it *Iterator) nextForSegment(segIdx int) *Entry {
+	it.cache.locks[segIdx].Lock()
+	defer it.cache.locks[segIdx].Unlock()
+	seg := &it.cache.segments[segIdx]
+	for it.slotIdx < 256 {
+		entry := it.nextForSlot(seg, it.slotIdx)
+		if entry != nil {
+			return entry
+		}
+		it.slotIdx++
+		it.entryIdx = 0
+	}
+	return nil
+}
+
+func (it *Iterator) nextForSlot(seg *segment, slotId int) *Entry {
+	slotOff := int32(it.slotIdx) * seg.slotCap
+	slot := seg.slotsData[slotOff : slotOff+seg.slotLens[it.slotIdx] : slotOff+seg.slotCap]
+	for it.entryIdx < len(slot) {
+		ptr := slot[it.entryIdx]
+		it.entryIdx++
+		now := uint32(time.Now().Unix())
+		var hdrBuf [ENTRY_HDR_SIZE]byte
+		seg.rb.ReadAt(hdrBuf[:], ptr.offset)
+		hdr := (*entryHdr)(unsafe.Pointer(&hdrBuf[0]))
+		if hdr.expireAt == 0 || hdr.expireAt > now {
+			entry := new(Entry)
+			entry.Key = make([]byte, hdr.keyLen)
+			entry.Value = make([]byte, hdr.valLen)
+			seg.rb.ReadAt(entry.Key, ptr.offset+ENTRY_HDR_SIZE)
+			seg.rb.ReadAt(entry.Value, ptr.offset+ENTRY_HDR_SIZE+int64(hdr.keyLen))
+			return entry
+		}
+	}
+	return nil
+}
+
+// NewIterator creates a new iterator for the cache.
+func (cache *Cache) NewIterator() *Iterator {
+	return &Iterator{
+		cache: cache,
+	}
+}