github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/swarmkit/watch/queue/queue_test.go (about)

     1  package queue
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/docker/go-events"
    10  	"github.com/sirupsen/logrus"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  type mockSink struct {
    15  	closed   bool
    16  	holdChan chan struct{}
    17  	data     []events.Event
    18  	mutex    sync.Mutex
    19  	once     sync.Once
    20  }
    21  
    22  func (s *mockSink) Write(event events.Event) error {
    23  	<-s.holdChan
    24  
    25  	s.mutex.Lock()
    26  	defer s.mutex.Unlock()
    27  	if s.closed {
    28  		return events.ErrSinkClosed
    29  	}
    30  	s.data = append(s.data, event)
    31  	return nil
    32  }
    33  
    34  func (s *mockSink) Close() error {
    35  	s.mutex.Lock()
    36  	defer s.mutex.Unlock()
    37  
    38  	s.once.Do(func() {
    39  		s.closed = true
    40  		close(s.holdChan)
    41  	})
    42  	return nil
    43  }
    44  
    45  func (s *mockSink) Len() int {
    46  	s.mutex.Lock()
    47  	defer s.mutex.Unlock()
    48  	return len(s.data)
    49  }
    50  
    51  func (s *mockSink) String() string {
    52  	s.mutex.Lock()
    53  	defer s.mutex.Unlock()
    54  	return fmt.Sprintf("%v", s.data)
    55  }
    56  
    57  func TestLimitQueueNoLimit(t *testing.T) {
    58  	require := require.New(t)
    59  	ch := make(chan struct{})
    60  	ms := &mockSink{
    61  		holdChan: ch,
    62  	}
    63  
    64  	// Create a limit queue with no limit and store 10k events. The events
    65  	// should be held in the queue until we unblock the sink.
    66  	q := NewLimitQueue(ms, 0)
    67  	defer q.Close()
    68  	defer ms.Close()
    69  
    70  	// Writing one event to the queue should block during the sink write phase
    71  	require.NoError(q.Write("test event"))
    72  
    73  	// Make sure the consumer goroutine receives the event
    74  	deadline := time.Now().Add(5 * time.Second)
    75  	for time.Now().Before(deadline) && q.Len() != 0 {
    76  		time.Sleep(20 * time.Millisecond)
    77  	}
    78  	require.Equal(0, q.Len())
    79  	require.Equal(0, ms.Len())
    80  
    81  	for i := 0; i < 9999; i++ {
    82  		require.NoError(q.Write("test event"))
    83  	}
    84  	require.Equal(9999, q.Len()) // 1 event blocked in the sink, 9999 waiting in the queue
    85  	require.Equal(0, ms.Len())
    86  
    87  	// Unblock the sink and expect all the events to have been flushed out of
    88  	// the queue.
    89  	for i := 0; i < 10000; i++ {
    90  		ch <- struct{}{}
    91  	}
    92  	deadline = time.Now().Add(5 * time.Second)
    93  	for time.Now().Before(deadline) && ms.Len() != 10000 {
    94  		time.Sleep(20 * time.Millisecond)
    95  	}
    96  
    97  	require.Equal(0, q.Len())
    98  	require.Equal(10000, ms.Len())
    99  }
   100  
   101  // TestLimitQueueWithLimit ensures that the limit queue works with a limit.
   102  func TestLimitQueueWithLimit(t *testing.T) {
   103  	require := require.New(t)
   104  	ch := make(chan struct{})
   105  	ms := &mockSink{
   106  		holdChan: ch,
   107  	}
   108  
   109  	// Create a limit queue with no limit and store 10k events. The events should be held in
   110  	// the queue until we unblock the sink.
   111  	q := NewLimitQueue(ms, 10)
   112  	defer q.Close()
   113  	defer ms.Close()
   114  
   115  	// Write the first event and wait for it to block on the writer
   116  	require.NoError(q.Write("test event"))
   117  	deadline := time.Now().Add(5 * time.Second)
   118  	for time.Now().Before(deadline) && q.Len() != 0 {
   119  		time.Sleep(20 * time.Millisecond)
   120  	}
   121  	require.Equal(0, ms.Len())
   122  	require.Equal(0, q.Len())
   123  
   124  	// Fill up the queue
   125  	for i := 0; i < 10; i++ {
   126  		require.NoError(q.Write("test event"))
   127  	}
   128  	require.Equal(0, ms.Len())
   129  	require.Equal(10, q.Len())
   130  
   131  	// Reading one event by the sink should allow us to write one more back
   132  	// without closing the queue.
   133  	ch <- struct{}{}
   134  	deadline = time.Now().Add(5 * time.Second)
   135  	for time.Now().Before(deadline) && q.Len() != 9 {
   136  		time.Sleep(20 * time.Millisecond)
   137  	}
   138  	require.Equal(9, q.Len())
   139  	require.Equal(1, ms.Len())
   140  	require.NoError(q.Write("test event"))
   141  	require.Equal(10, q.Len())
   142  	require.Equal(1, ms.Len())
   143  
   144  	// Trying to write a new event in the queue should flush it
   145  	logrus.Debugf("Closing queue")
   146  	err := q.Write("test event")
   147  	require.Error(err)
   148  	require.Equal(ErrQueueFull, err)
   149  	require.Equal(10, q.Len())
   150  	require.Equal(1, ms.Len())
   151  
   152  	// Further writes should return the same error
   153  	err = q.Write("test event")
   154  	require.Error(err)
   155  	require.Equal(ErrQueueFull, err)
   156  	require.Equal(10, q.Len())
   157  	require.Equal(1, ms.Len())
   158  
   159  	// Reading one event from the sink will allow one more write to go through again
   160  	ch <- struct{}{}
   161  	deadline = time.Now().Add(5 * time.Second)
   162  	for time.Now().Before(deadline) && q.Len() != 9 {
   163  		time.Sleep(20 * time.Millisecond)
   164  	}
   165  	require.Equal(9, q.Len())
   166  	require.Equal(2, ms.Len())
   167  	require.NoError(q.Write("test event"))
   168  	require.Equal(10, q.Len())
   169  	require.Equal(2, ms.Len())
   170  
   171  	err = q.Write("test event")
   172  	require.Error(err)
   173  	require.Equal(ErrQueueFull, err)
   174  	require.Equal(10, q.Len())
   175  	require.Equal(2, ms.Len())
   176  }