diff --git a/mruinvmap.go b/mruinvmap.go index f58648c1..3fa0c1c6 100644 --- a/mruinvmap.go +++ b/mruinvmap.go @@ -5,16 +5,17 @@ package main import ( + "container/list" "fmt" "github.com/conformal/btcwire" - "time" ) // MruInventoryMap provides a map that is limited to a maximum number of items // with eviction for the oldest entry when the limit is exceeded. type MruInventoryMap struct { - invMap map[btcwire.InvVect]int64 // Use int64 for time for less mem. - limit uint + invMap map[btcwire.InvVect]*list.Element // nearly O(1) lookups + invList *list.List // O(1) insert, update, delete + limit uint } // String returns the map as a human-readable string. @@ -34,39 +35,51 @@ func (m *MruInventoryMap) Exists(iv *btcwire.InvVect) bool { // item if adding the new item would exceed the max limit. func (m *MruInventoryMap) Add(iv *btcwire.InvVect) { // When the limit is zero, nothing can be added to the map, so just - // return + // return. if m.limit == 0 { return } - // When the entry already exists update its last seen time. - if m.Exists(iv) { - m.invMap[*iv] = time.Now().Unix() + // When the entry already exists move it to the front of the list + // thereby marking it most recently used. + if node, exists := m.invMap[*iv]; exists { + m.invList.MoveToFront(node) return } - // Evict the oldest entry if the the new entry would exceed the size - // limit for the map. + // Evict the least recently used entry (back of the list) if the the new + // entry would exceed the size limit for the map. Also reuse the list + // node so a new one doesn't have to be allocated. if uint(len(m.invMap))+1 > m.limit { - var oldestEntry btcwire.InvVect - var oldestTime int64 - for iv, lastUpdated := range m.invMap { - if oldestTime == 0 || lastUpdated < oldestTime { - oldestEntry = iv - oldestTime = lastUpdated - } + node := m.invList.Back() + lru, ok := node.Value.(*btcwire.InvVect) + if !ok { + return } - m.Delete(&oldestEntry) + // Evict least recently used item. + delete(m.invMap, *lru) + + // Reuse the list node of the item that was just evicted for the + // new item. + node.Value = iv + m.invList.MoveToFront(node) + m.invMap[*iv] = node + return } - m.invMap[*iv] = time.Now().Unix() + // The limit hasn't been reached yet, so just add the new item. + node := m.invList.PushFront(iv) + m.invMap[*iv] = node return } // Delete deletes the passed inventory item from the map (if it exists). func (m *MruInventoryMap) Delete(iv *btcwire.InvVect) { - delete(m.invMap, *iv) + if node, exists := m.invMap[*iv]; exists { + m.invList.Remove(node) + delete(m.invMap, *iv) + } } // NewMruInventoryMap returns a new inventory map that is limited to the number @@ -75,8 +88,9 @@ func (m *MruInventoryMap) Delete(iv *btcwire.InvVect) { // new entry.. func NewMruInventoryMap(limit uint) *MruInventoryMap { m := MruInventoryMap{ - invMap: make(map[btcwire.InvVect]int64), - limit: limit, + invMap: make(map[btcwire.InvVect]*list.Element), + invList: list.New(), + limit: limit, } return &m }