github.com/weaviate/weaviate@v1.24.6/usecases/memwatch/monitor_test.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  	"runtime/debug"
    16  	"testing"
    17  
    18  	"github.com/stretchr/testify/assert"
    19  )
    20  
    21  func TestMonitor(t *testing.T) {
    22  	t.Run("with constant profiles (no changes)", func(t *testing.T) {
    23  		metrics := &fakeHeapReader{val: 30000}
    24  		limiter := &fakeLimitSetter{limit: 100000}
    25  
    26  		m := NewMonitor(metrics.Read, limiter.SetMemoryLimit, 0.97)
    27  		m.Refresh()
    28  
    29  		assert.Equal(t, 0.3, m.Ratio())
    30  	})
    31  
    32  	t.Run("with less memory than the threshold", func(t *testing.T) {
    33  		metrics := &fakeHeapReader{val: 700 * MiB}
    34  		limiter := &fakeLimitSetter{limit: 1 * GiB}
    35  
    36  		m := NewMonitor(metrics.Read, limiter.SetMemoryLimit, 0.97)
    37  		m.Refresh()
    38  
    39  		err := m.CheckAlloc(100 * MiB)
    40  		assert.NoError(t, err, "with 700 allocated, an additional 100 would be about 80% which is not a problem")
    41  
    42  		err = m.CheckAlloc(299 * MiB)
    43  		assert.Error(t, err, "with 700 allocated, an additional 299 would be about 97.5% which is not allowed")
    44  
    45  		err = m.CheckAlloc(400 * MiB)
    46  		assert.Error(t, err, "with 700 allocated, an additional 400 would be about 110% which is not allowed")
    47  	})
    48  
    49  	t.Run("with memory already over the threshold", func(t *testing.T) {
    50  		metrics := &fakeHeapReader{val: 1025 * MiB}
    51  		limiter := &fakeLimitSetter{limit: 1 * GiB}
    52  
    53  		m := NewMonitor(metrics.Read, limiter.SetMemoryLimit, 0.97)
    54  		m.Refresh()
    55  
    56  		err := m.CheckAlloc(1 * B)
    57  		assert.Error(t, err,
    58  			"any check should fail, since we're already over the limit")
    59  
    60  		err = m.CheckAlloc(10 * MiB)
    61  		assert.Error(t, err, "any check should fail, since we're already over the limit")
    62  
    63  		err = m.CheckAlloc(1 * TiB)
    64  		assert.Error(t, err, "any check should fail, since we're already over the limit")
    65  	})
    66  
    67  	t.Run("with real dependencies", func(t *testing.T) {
    68  		m := NewMonitor(LiveHeapReader, debug.SetMemoryLimit, 0.97)
    69  		_ = m.Ratio()
    70  	})
    71  }
    72  
    73  type fakeHeapReader struct {
    74  	val int64
    75  }
    76  
    77  func (f fakeHeapReader) Read() int64 {
    78  	return f.val
    79  }
    80  
    81  type fakeLimitSetter struct {
    82  	limit int64
    83  }
    84  
    85  func (f *fakeLimitSetter) SetMemoryLimit(newLimit int64) int64 {
    86  	if newLimit >= 0 {
    87  		panic("should have been read only")
    88  	}
    89  
    90  	return f.limit
    91  }