github.com/go-graphite/carbonapi@v0.17.0/zipper/cache/query.go (about)

     1  package cache
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"sync/atomic"
     7  
     8  	"github.com/dgryski/go-expirecache"
     9  )
    10  
    11  const (
    12  	Empty uint64 = iota
    13  	QueryIsPending
    14  	DataIsAvailable
    15  )
    16  
    17  type QueryItem struct {
    18  	sync.RWMutex
    19  	Key           string
    20  	Data          atomic.Value
    21  	Flags         uint64 // DataIsAvailable or QueryIsPending
    22  	QueryFinished chan struct{}
    23  
    24  	parent *QueryCache
    25  }
    26  
    27  func (q *QueryItem) GetStatus() uint64 {
    28  	s := atomic.LoadUint64(&q.Flags)
    29  	return s
    30  }
    31  
    32  func (q *QueryItem) FetchOrLock(ctx context.Context) (interface{}, bool) {
    33  	d := q.Data.Load()
    34  	if d != nil {
    35  		return d, true
    36  	}
    37  
    38  	ok := atomic.CompareAndSwapUint64(&q.Flags, Empty, QueryIsPending)
    39  	if ok {
    40  		// We are the leader now and will be fetching the data
    41  		return nil, false
    42  	}
    43  
    44  	q.RLock()
    45  	defer q.RUnlock()
    46  
    47  	select {
    48  	case <-ctx.Done():
    49  		return nil, true
    50  	case <-q.QueryFinished:
    51  		break
    52  	}
    53  
    54  	return q.Data.Load(), true
    55  }
    56  
    57  func (q *QueryItem) StoreAbort() {
    58  	d := q.Data.Load()
    59  	if d != nil {
    60  		return
    61  	}
    62  	atomic.StoreUint64(&q.Flags, Empty)
    63  	close(q.QueryFinished)
    64  
    65  	q.Lock()
    66  	q.QueryFinished = make(chan struct{})
    67  	q.Unlock()
    68  }
    69  
    70  func (q *QueryItem) StoreAndUnlock(data interface{}, size uint64) {
    71  	q.Data.Store(data)
    72  	atomic.StoreUint64(&q.Flags, DataIsAvailable)
    73  	close(q.QueryFinished)
    74  	atomic.AddUint64(&q.parent.totalSize, size)
    75  }
    76  
    77  type QueryCache struct {
    78  	ec *expirecache.Cache
    79  
    80  	objectCount uint64
    81  	totalSize   uint64
    82  	expireTime  int32
    83  }
    84  
    85  func NewQueryCache(queryCacheSizeMB uint64, expireTime int32) *QueryCache {
    86  	return &QueryCache{
    87  		ec:         expirecache.New(queryCacheSizeMB),
    88  		expireTime: expireTime,
    89  	}
    90  }
    91  
    92  // TODO: Make size and expire configurable
    93  func (q *QueryCache) GetQueryItem(k string) *QueryItem {
    94  	objectCount := atomic.AddUint64(&q.objectCount, 1)
    95  	size := atomic.AddUint64(&q.totalSize, 1)
    96  	emptyQueryItem := &QueryItem{
    97  		Key:           k,
    98  		QueryFinished: make(chan struct{}),
    99  		Flags:         Empty,
   100  
   101  		parent: q,
   102  	}
   103  	item := q.ec.GetOrSet(k, emptyQueryItem, size/objectCount, q.expireTime).(*QueryItem)
   104  
   105  	return item
   106  }