From 494c19ec2b41d75a66777c1494064f0b70e121dc Mon Sep 17 00:00:00 2001
From: David Skorvaga <55138449+pheepi@users.noreply.github.com>
Date: Tue, 14 Jan 2020 20:21:46 +0100
Subject: [PATCH] Add custom timer interface (#69)

---
 cache.go    | 10 +++++++++-
 iterator.go |  3 +--
 segment.go  | 22 ++++++++++++++++++----
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/cache.go b/cache.go
index 9dedbf8..644de36 100644
--- a/cache.go
+++ b/cache.go
@@ -32,12 +32,20 @@ func hashFunc(data []byte) uint64 {
 // `debug.SetGCPercent()`, set it to a much smaller value
 // to limit the memory consumption and GC pause time.
 func NewCache(size int) (cache *Cache) {
+	return NewCacheCustomTimer(size, defaultTimer{})
+}
+
+// NewCacheCustomTimer returns new cache with custom timer.
+func NewCacheCustomTimer(size int, timer Timer) (cache *Cache) {
 	if size < minBufSize {
 		size = minBufSize
 	}
+	if timer == nil {
+		timer = defaultTimer{}
+	}
 	cache = new(Cache)
 	for i := 0; i < segmentCount; i++ {
-		cache.segments[i] = newSegment(size/segmentCount, i)
+		cache.segments[i] = newSegment(size/segmentCount, i, timer)
 	}
 	return
 }
diff --git a/iterator.go b/iterator.go
index de07800..6f83d99 100644
--- a/iterator.go
+++ b/iterator.go
@@ -1,7 +1,6 @@
 package freecache
 
 import (
-	"time"
 	"unsafe"
 )
 
@@ -56,7 +55,7 @@ func (it *Iterator) nextForSlot(seg *segment, slotId int) *Entry {
 	for it.entryIdx < len(slot) {
 		ptr := slot[it.entryIdx]
 		it.entryIdx++
-		now := uint32(time.Now().Unix())
+		now := seg.timer.Now()
 		var hdrBuf [ENTRY_HDR_SIZE]byte
 		seg.rb.ReadAt(hdrBuf[:], ptr.offset)
 		hdr := (*entryHdr)(unsafe.Pointer(&hdrBuf[0]))
diff --git a/segment.go b/segment.go
index cb996f6..9409531 100644
--- a/segment.go
+++ b/segment.go
@@ -14,6 +14,12 @@ var ErrLargeKey = errors.New("The key is larger than 65535")
 var ErrLargeEntry = errors.New("The entry size is larger than 1/1024 of cache size")
 var ErrNotFound = errors.New("Entry not found")
 
+// timer holds representation of current time.
+type Timer interface {
+	// Give current time (in seconds)
+	Now() uint32
+}
+
 // entry pointer struct points to an entry in ring buffer
 type entryPtr struct {
 	offset   int64  // entry offset in ring buffer
@@ -46,6 +52,7 @@ type segment struct {
 	entryCount    int64
 	totalCount    int64      // number of entries in ring buffer, including deleted entries.
 	totalTime     int64      // used to calculate least recent used entry.
+	timer         Timer      // Timer giving current time
 	totalEvacuate int64      // used for debug
 	totalExpired  int64      // used for debug
 	overwrites    int64      // used for debug
@@ -55,9 +62,16 @@ type segment struct {
 	slotsData     []entryPtr // shared by all 256 slots
 }
 
-func newSegment(bufSize int, segId int) (seg segment) {
+type defaultTimer struct{}
+
+func (timer defaultTimer) Now() uint32 {
+	return uint32(time.Now().Unix())
+}
+
+func newSegment(bufSize int, segId int, timer Timer) (seg segment) {
 	seg.rb = NewRingBuf(bufSize, 0)
 	seg.segId = segId
+	seg.timer = timer
 	seg.vacuumLen = int64(bufSize)
 	seg.slotCap = 1
 	seg.slotsData = make([]entryPtr, 256*seg.slotCap)
@@ -73,7 +87,7 @@ func (seg *segment) set(key, value []byte, hashVal uint64, expireSeconds int) (e
 		// Do not accept large entry.
 		return ErrLargeEntry
 	}
-	now := uint32(time.Now().Unix())
+	now := seg.timer.Now()
 	expireAt := uint32(0)
 	if expireSeconds > 0 {
 		expireAt = now + uint32(expireSeconds)
@@ -207,7 +221,7 @@ func (seg *segment) get(key, buf []byte, hashVal uint64, peek bool) (value []byt
 	seg.rb.ReadAt(hdrBuf[:], ptr.offset)
 	hdr := (*entryHdr)(unsafe.Pointer(&hdrBuf[0]))
 	if !peek {
-		now := uint32(time.Now().Unix())
+		now := seg.timer.Now()
 		expireAt = hdr.expireAt
 
 		if hdr.expireAt != 0 && hdr.expireAt <= now {
@@ -256,7 +270,7 @@ func (seg *segment) ttl(key []byte, hashVal uint64) (timeLeft uint32, err error)
 		return
 	}
 	ptr := &slot[idx]
-	now := uint32(time.Now().Unix())
+	now := seg.timer.Now()
 
 	var hdrBuf [ENTRY_HDR_SIZE]byte
 	seg.rb.ReadAt(hdrBuf[:], ptr.offset)