github.com/bodgit/sevenzip@v1.5.1/internal/pool/pool.go (about)

     1  package pool
     2  
     3  import (
     4  	"container/list"
     5  	"runtime"
     6  	"sort"
     7  	"sync"
     8  
     9  	"github.com/bodgit/sevenzip/internal/util"
    10  )
    11  
    12  // Pooler is the interface implemented by a pool.
    13  type Pooler interface {
    14  	Get(int64) (util.SizeReadSeekCloser, bool)
    15  	Put(int64, util.SizeReadSeekCloser) (bool, error)
    16  }
    17  
    18  // Constructor is the function prototype used to instantiate a pool.
    19  type Constructor func() (Pooler, error)
    20  
    21  type noopPool struct{}
    22  
    23  // NewNoopPool returns a Pooler that doesn't actually pool anything.
    24  func NewNoopPool() (Pooler, error) {
    25  	return new(noopPool), nil
    26  }
    27  
    28  func (noopPool) Get(_ int64) (util.SizeReadSeekCloser, bool) {
    29  	return nil, false
    30  }
    31  
    32  func (noopPool) Put(_ int64, rc util.SizeReadSeekCloser) (bool, error) {
    33  	return false, rc.Close()
    34  }
    35  
    36  type pool struct {
    37  	mutex     sync.Mutex
    38  	size      int
    39  	evictList *list.List
    40  	items     map[int64]*list.Element
    41  }
    42  
    43  type entry struct {
    44  	key   int64
    45  	value util.SizeReadSeekCloser
    46  }
    47  
    48  // NewPool returns a Pooler that uses a LRU strategy to maintain a fixed pool
    49  // of util.SizeReadSeekCloser's keyed by their stream offset.
    50  func NewPool() (Pooler, error) {
    51  	return &pool{
    52  		size:      runtime.NumCPU(),
    53  		evictList: list.New(),
    54  		items:     make(map[int64]*list.Element),
    55  	}, nil
    56  }
    57  
    58  func (p *pool) Get(offset int64) (util.SizeReadSeekCloser, bool) {
    59  	p.mutex.Lock()
    60  	defer p.mutex.Unlock()
    61  
    62  	if ent, ok := p.items[offset]; ok {
    63  		_ = p.removeElement(ent, false)
    64  
    65  		return ent.Value.(*entry).value, true //nolint:forcetypeassert
    66  	}
    67  
    68  	// Sort keys in descending order
    69  	keys := p.keys()
    70  	sort.Slice(keys, func(i, j int) bool { return keys[i] > keys[j] })
    71  
    72  	for _, k := range keys {
    73  		// First key less than offset is the closest
    74  		if k < offset {
    75  			ent := p.items[k]
    76  			_ = p.removeElement(ent, false)
    77  
    78  			return ent.Value.(*entry).value, true //nolint:forcetypeassert
    79  		}
    80  	}
    81  
    82  	return nil, false
    83  }
    84  
    85  func (p *pool) Put(offset int64, rc util.SizeReadSeekCloser) (bool, error) {
    86  	p.mutex.Lock()
    87  	defer p.mutex.Unlock()
    88  
    89  	if _, ok := p.items[offset]; ok {
    90  		return false, nil
    91  	}
    92  
    93  	ent := &entry{offset, rc}
    94  	entry := p.evictList.PushFront(ent)
    95  	p.items[offset] = entry
    96  
    97  	var err error
    98  
    99  	evict := p.evictList.Len() > p.size
   100  	if evict {
   101  		err = p.removeOldest()
   102  	}
   103  
   104  	return evict, err
   105  }
   106  
   107  func (p *pool) keys() []int64 {
   108  	keys := make([]int64, len(p.items))
   109  	i := 0
   110  
   111  	for ent := p.evictList.Back(); ent != nil; ent = ent.Prev() {
   112  		keys[i] = ent.Value.(*entry).key //nolint:forcetypeassert
   113  		i++
   114  	}
   115  
   116  	return keys
   117  }
   118  
   119  func (p *pool) removeOldest() error {
   120  	if ent := p.evictList.Back(); ent != nil {
   121  		return p.removeElement(ent, true)
   122  	}
   123  
   124  	return nil
   125  }
   126  
   127  func (p *pool) removeElement(e *list.Element, cb bool) error {
   128  	p.evictList.Remove(e)
   129  	kv := e.Value.(*entry) //nolint:forcetypeassert
   130  	delete(p.items, kv.key)
   131  
   132  	if cb {
   133  		return kv.value.Close()
   134  	}
   135  
   136  	return nil
   137  }