1 files changed, 32 insertions, 8 deletions
diff --git a/vendor/go.opencensus.io/trace/lrumap.go b/vendor/go.opencensus.io/trace/lrumap.go
index 3f80a3368..dc7a295c7 100644
--- a/vendor/go.opencensus.io/trace/lrumap.go
+++ b/vendor/go.opencensus.io/trace/lrumap.go
@@ -15,23 +15,47 @@
 package trace
 
 import (
-	"github.com/hashicorp/golang-lru/simplelru"
+	"github.com/golang/groupcache/lru"
 )
 
+// A simple lru.Cache wrapper that tracks the keys of the current contents and
+// the cumulative number of evicted items.
 type lruMap struct {
-	simpleLruMap *simplelru.LRU
+	cacheKeys    map[lru.Key]bool
+	cache        *lru.Cache
 	droppedCount int
 }
 
 func newLruMap(size int) *lruMap {
-	lm := &lruMap{}
-	lm.simpleLruMap, _ = simplelru.NewLRU(size, nil)
+	lm := &lruMap{
+		cacheKeys:    make(map[lru.Key]bool),
+		cache:        lru.New(size),
+		droppedCount: 0,
+	}
+	lm.cache.OnEvicted = func(key lru.Key, value interface{}) {
+		delete(lm.cacheKeys, key)
+		lm.droppedCount++
+	}
 	return lm
 }
 
-func (lm *lruMap) add(key, value interface{}) {
-	evicted := lm.simpleLruMap.Add(key, value)
-	if evicted {
-		lm.droppedCount++
+func (lm lruMap) len() int {
+	return lm.cache.Len()
+}
+
+func (lm lruMap) keys() []interface{} {
+	keys := []interface{}{}
+	for k := range lm.cacheKeys {
+		keys = append(keys, k)
 	}
+	return keys
+}
+
+func (lm *lruMap) add(key, value interface{}) {
+	lm.cacheKeys[lru.Key(key)] = true
+	lm.cache.Add(lru.Key(key), value)
+}
+
+func (lm *lruMap) get(key interface{}) (interface{}, bool) {
+	return lm.cache.Get(key)
 }