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

     1  package ingester
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  type dumbFlusher struct {
    12  	onFlush func()
    13  }
    14  
    15  func newDumbFlusher(onFlush func()) *dumbFlusher {
    16  	return &dumbFlusher{
    17  		onFlush: onFlush,
    18  	}
    19  }
    20  
    21  func (f *dumbFlusher) Flush() {
    22  	if f.onFlush != nil {
    23  		f.onFlush()
    24  	}
    25  }
    26  
    27  func nilMetrics() *ingesterMetrics { return newIngesterMetrics(nil) }
    28  
    29  func TestReplayController(t *testing.T) {
    30  	var ops []string
    31  	var opLock sync.Mutex
    32  
    33  	var rc *replayController
    34  	flusher := newDumbFlusher(
    35  		func() {
    36  			rc.Sub(100) // simulate flushing 100 bytes
    37  			opLock.Lock()
    38  			defer opLock.Unlock()
    39  			ops = append(ops, "Flush")
    40  		},
    41  	)
    42  	rc = newReplayController(nilMetrics(), WALConfig{ReplayMemoryCeiling: 100}, flusher)
    43  
    44  	var wg sync.WaitGroup
    45  	n := 5
    46  	wg.Add(n)
    47  
    48  	for i := 0; i < n; i++ {
    49  		// In order to prevent all the goroutines from running before they've added bytes
    50  		// to the internal count, introduce a brief sleep.
    51  		time.Sleep(time.Millisecond)
    52  
    53  		// nolint:errcheck,unparam
    54  		go rc.WithBackPressure(func() error {
    55  			rc.Add(50)
    56  			opLock.Lock()
    57  			defer opLock.Unlock()
    58  			ops = append(ops, "WithBackPressure")
    59  			wg.Done()
    60  			return nil
    61  		})
    62  	}
    63  
    64  	wg.Wait()
    65  
    66  	expected := []string{
    67  		"WithBackPressure", // add 50, total 50
    68  		"WithBackPressure", // add 50, total 100
    69  		"Flush",            // subtract 100, total 0
    70  		"WithBackPressure", // add 50, total 50
    71  		"WithBackPressure", // add 50, total 100
    72  		"Flush",            // subtract 100, total 0
    73  		"WithBackPressure", // add 50, total 50
    74  	}
    75  	require.Equal(t, expected, ops)
    76  
    77  }