github.com/weaviate/weaviate@v1.24.6/usecases/memwatch/monitor.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package memwatch
    13  
    14  import (
    15  	"fmt"
    16  	"runtime/metrics"
    17  	"sync"
    18  )
    19  
    20  const (
    21  	B   = 1
    22  	KiB = 1 << (10 * iota) // 2^10
    23  	MiB = 1 << (10 * iota) // 2^20
    24  	GiB = 1 << (10 * iota) // 2^30
    25  	TiB = 1 << (10 * iota) // 2^40
    26  )
    27  
    28  var ErrNotEnoughMemory = fmt.Errorf("not enough memory")
    29  
    30  // Monitor allows making statements about the memory ratio used by the application
    31  type Monitor struct {
    32  	metricsReader metricsReader
    33  	limitSetter   limitSetter
    34  	maxRatio      float64
    35  
    36  	// state
    37  	mu    sync.Mutex
    38  	limit int64
    39  	used  int64
    40  }
    41  
    42  // Refresh retrieves the current memory stats from the runtime and stores them
    43  // in the local cache
    44  func (m *Monitor) Refresh() {
    45  	m.obtainCurrentUsage()
    46  	m.updateLimit()
    47  }
    48  
    49  // we have no intentions of ever modifying the limit, but SetMemoryLimit with a
    50  // negative value is the only way to read the limit from the runtime
    51  type limitSetter func(size int64) int64
    52  
    53  // NewMonitor creates a [Monitor] with the given metrics reader and target
    54  // ratio
    55  //
    56  // Typically this would be called with LiveHeapReader and
    57  // debug.SetMemoryLimit
    58  func NewMonitor(metricsReader metricsReader, limitSetter limitSetter,
    59  	maxRatio float64,
    60  ) *Monitor {
    61  	return &Monitor{
    62  		metricsReader: metricsReader,
    63  		limitSetter:   limitSetter,
    64  		maxRatio:      maxRatio,
    65  	}
    66  }
    67  
    68  func (m *Monitor) CheckAlloc(sizeInBytes int64) error {
    69  	m.mu.Lock()
    70  	defer m.mu.Unlock()
    71  
    72  	if float64(m.used+sizeInBytes)/float64(m.limit) > m.maxRatio {
    73  		return ErrNotEnoughMemory
    74  	}
    75  
    76  	return nil
    77  }
    78  
    79  func (m *Monitor) Ratio() float64 {
    80  	m.mu.Lock()
    81  	defer m.mu.Unlock()
    82  
    83  	return float64(m.used) / float64(m.limit)
    84  }
    85  
    86  // obtainCurrentUsage obtains the most recent live heap from runtime/metrics
    87  func (m *Monitor) obtainCurrentUsage() {
    88  	m.setUsed(m.metricsReader())
    89  }
    90  
    91  func LiveHeapReader() int64 {
    92  	const liveHeapBytesMetric = "/gc/heap/live:bytes"
    93  	sample := make([]metrics.Sample, 1)
    94  	sample[0].Name = liveHeapBytesMetric
    95  	metrics.Read(sample)
    96  
    97  	if sample[0].Value.Kind() == metrics.KindBad {
    98  		panic(fmt.Sprintf("metric %q no longer supported", liveHeapBytesMetric))
    99  	}
   100  
   101  	return int64(sample[0].Value.Uint64())
   102  }
   103  
   104  // setUsed is a thread-safe way way to set the current usage
   105  func (m *Monitor) setUsed(used int64) {
   106  	m.mu.Lock()
   107  	defer m.mu.Unlock()
   108  
   109  	m.used = used
   110  }
   111  
   112  func (m *Monitor) updateLimit() {
   113  	m.mu.Lock()
   114  	defer m.mu.Unlock()
   115  
   116  	// setting a negative limit is the only way to obtain the current limit
   117  	m.limit = m.limitSetter(-1)
   118  }
   119  
   120  type metricsReader func() int64