github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/ingester/streams_map.go (about)

     1  package ingester
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/prometheus/common/model"
     7  	"go.uber.org/atomic"
     8  )
     9  
    10  type streamsMap struct {
    11  	consistencyMtx sync.RWMutex // Keep read/write consistency between other fields
    12  	streams        *sync.Map    // map[string]*stream
    13  	streamsByFP    *sync.Map    // map[model.Fingerprint]*stream
    14  	streamsCounter *atomic.Int64
    15  }
    16  
    17  func newStreamsMap() *streamsMap {
    18  	return &streamsMap{
    19  		consistencyMtx: sync.RWMutex{},
    20  		streams:        &sync.Map{},
    21  		streamsByFP:    &sync.Map{},
    22  		streamsCounter: atomic.NewInt64(0),
    23  	}
    24  }
    25  
    26  // Load is lock-free. If usage of the stream is consistency sensitive, must be called inside WithRLock at least
    27  func (m *streamsMap) Load(key string) (*stream, bool) {
    28  	return m.load(m.streams, key)
    29  }
    30  
    31  // LoadByFP is lock-free. If usage of the stream is consistency sensitive, must be called inside WithRLock at least
    32  func (m *streamsMap) LoadByFP(fp model.Fingerprint) (*stream, bool) {
    33  	return m.load(m.streamsByFP, fp)
    34  }
    35  
    36  // Store must be called inside WithLock
    37  func (m *streamsMap) Store(key string, s *stream) {
    38  	m.store(key, s)
    39  }
    40  
    41  // StoreByFP must be called inside WithLock
    42  func (m *streamsMap) StoreByFP(fp model.Fingerprint, s *stream) {
    43  	m.store(fp, s)
    44  }
    45  
    46  // Delete must be called inside WithLock
    47  func (m *streamsMap) Delete(s *stream) bool {
    48  	_, loaded := m.streams.LoadAndDelete(s.labelsString)
    49  	if loaded {
    50  		m.streamsByFP.Delete(s.fp)
    51  		m.streamsCounter.Dec()
    52  		return true
    53  	}
    54  	return false
    55  }
    56  
    57  // LoadOrStoreNew already has lock inside, do NOT call inside WithLock or WithRLock
    58  func (m *streamsMap) LoadOrStoreNew(key string, newStreamFn func() (*stream, error), postLoadFn func(*stream) error) (*stream, bool, error) {
    59  	return m.loadOrStoreNew(m.streams, key, newStreamFn, postLoadFn)
    60  }
    61  
    62  // LoadOrStoreNewByFP already has lock inside, do NOT call inside WithLock or WithRLock
    63  func (m *streamsMap) LoadOrStoreNewByFP(fp model.Fingerprint, newStreamFn func() (*stream, error), postLoadFn func(*stream) error) (*stream, bool, error) {
    64  	return m.loadOrStoreNew(m.streamsByFP, fp, newStreamFn, postLoadFn)
    65  }
    66  
    67  // WithLock is a helper function to execute write operations
    68  func (m *streamsMap) WithLock(fn func()) {
    69  	m.consistencyMtx.Lock()
    70  	defer m.consistencyMtx.Unlock()
    71  	fn()
    72  }
    73  
    74  // WithRLock is a helper function to execute consistency sensitive read operations.
    75  // Generally, if a stream loaded from streamsMap will have its chunkMtx locked, chunkMtx.Lock is supposed to be called
    76  // within this function.
    77  func (m *streamsMap) WithRLock(fn func()) {
    78  	m.consistencyMtx.RLock()
    79  	defer m.consistencyMtx.RUnlock()
    80  	fn()
    81  }
    82  
    83  func (m *streamsMap) ForEach(fn func(s *stream) (bool, error)) error {
    84  	var c bool
    85  	var err error
    86  	m.streams.Range(func(key, value interface{}) bool {
    87  		c, err = fn(value.(*stream))
    88  		return c
    89  	})
    90  	return err
    91  }
    92  
    93  func (m *streamsMap) Len() int {
    94  	return int(m.streamsCounter.Load())
    95  }
    96  
    97  func (m *streamsMap) load(mp *sync.Map, key interface{}) (*stream, bool) {
    98  	if v, ok := mp.Load(key); ok {
    99  		return v.(*stream), true
   100  	}
   101  	return nil, false
   102  }
   103  
   104  func (m *streamsMap) store(key interface{}, s *stream) {
   105  	if labelsString, ok := key.(string); ok {
   106  		m.streams.Store(labelsString, s)
   107  	} else {
   108  		m.streams.Store(s.labelsString, s)
   109  	}
   110  	m.streamsByFP.Store(s.fp, s)
   111  	m.streamsCounter.Inc()
   112  }
   113  
   114  // newStreamFn: Called if not loaded, with consistencyMtx locked. Must not be nil
   115  // postLoadFn: Called if loaded, with consistencyMtx read-locked at least. Can be nil
   116  func (m *streamsMap) loadOrStoreNew(mp *sync.Map, key interface{}, newStreamFn func() (*stream, error), postLoadFn func(*stream) error) (*stream, bool, error) {
   117  	var s *stream
   118  	var loaded bool
   119  	var err error
   120  	m.WithRLock(func() {
   121  		if s, loaded = m.load(mp, key); loaded {
   122  			if postLoadFn != nil {
   123  				err = postLoadFn(s)
   124  			}
   125  		}
   126  	})
   127  
   128  	if loaded {
   129  		return s, true, err
   130  	}
   131  
   132  	m.WithLock(func() {
   133  		// Double check
   134  		if s, loaded = m.load(mp, key); loaded {
   135  			if postLoadFn != nil {
   136  				err = postLoadFn(s)
   137  			}
   138  			return
   139  		}
   140  
   141  		s, err = newStreamFn()
   142  		if err != nil {
   143  			return
   144  		}
   145  		m.store(key, s)
   146  	})
   147  
   148  	return s, loaded, err
   149  }