github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/utils/async/ring_buffer_test.go (about)

     1  // Copyright 2021 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package async
    16  
    17  import (
    18  	"fmt"
    19  	"io"
    20  	"sync"
    21  	"testing"
    22  
    23  	"github.com/stretchr/testify/assert"
    24  )
    25  
    26  func TestSingleThread(t *testing.T) {
    27  	tests := []struct {
    28  		allocSize int
    29  		numItems  int
    30  	}{
    31  		{128, 127},
    32  		{128, 128},
    33  		{128, 129},
    34  		{1, 1024},
    35  	}
    36  
    37  	for _, test := range tests {
    38  		t.Run(fmt.Sprintf("alloc %d items %d", test.allocSize, test.numItems), func(t *testing.T) {
    39  			rb := NewRingBuffer(test.allocSize)
    40  
    41  			for i := 0; i < test.numItems; i++ {
    42  				err := rb.Push(i, rb.epoch)
    43  				assert.NoError(t, err)
    44  			}
    45  
    46  			for i := 0; i < test.numItems; i++ {
    47  				item, err := rb.Pop()
    48  				assert.NoError(t, err)
    49  				assert.Equal(t, i, item.(int))
    50  			}
    51  
    52  			item, ok := rb.TryPop()
    53  			assert.Nil(t, item)
    54  			assert.False(t, ok)
    55  		})
    56  	}
    57  }
    58  
    59  func TestOneProducerOneConsumer(t *testing.T) {
    60  	tests := []struct {
    61  		allocSize int
    62  		numItems  int
    63  	}{
    64  		{128, 127},
    65  		{128, 128},
    66  		{128, 129},
    67  		{1, 1024},
    68  	}
    69  
    70  	for _, test := range tests {
    71  		t.Run(fmt.Sprintf("alloc %d items %d", test.allocSize, test.numItems), func(t *testing.T) {
    72  			rb := NewRingBuffer(test.allocSize)
    73  
    74  			go func() {
    75  				defer rb.Close()
    76  
    77  				for i := 0; i < test.numItems; i++ {
    78  					err := rb.Push(i, rb.epoch)
    79  					assert.NoError(t, err)
    80  				}
    81  			}()
    82  
    83  			for i := 0; i < test.numItems; i++ {
    84  				item, err := rb.Pop()
    85  				assert.NoError(t, err)
    86  				assert.Equal(t, i, item.(int))
    87  			}
    88  
    89  			item, err := rb.Pop()
    90  			assert.Nil(t, item)
    91  			assert.Equal(t, io.EOF, err)
    92  		})
    93  	}
    94  }
    95  
    96  func TestNProducersNConsumers(t *testing.T) {
    97  	tests := []struct {
    98  		producers        int
    99  		consumers        int
   100  		allocSize        int
   101  		itemsPerProducer int
   102  	}{
   103  		{2, 8, 128, 127},
   104  		{2, 8, 128, 128},
   105  		{2, 8, 128, 129},
   106  		{2, 8, 1, 1024},
   107  		{8, 2, 1, 1024},
   108  		{8, 8, 1, 1024},
   109  	}
   110  
   111  	for _, test := range tests {
   112  		t.Run(fmt.Sprintf("producers %d consumers %d alloc %d items per producer %d", test.producers, test.consumers, test.allocSize, test.itemsPerProducer), func(t *testing.T) {
   113  			rb := NewRingBuffer(test.allocSize)
   114  
   115  			producerGroup := &sync.WaitGroup{}
   116  			producerGroup.Add(test.producers)
   117  			for i := 0; i < test.producers; i++ {
   118  				go func() {
   119  					defer producerGroup.Done()
   120  					for i := 0; i < test.itemsPerProducer; i++ {
   121  						err := rb.Push(i, rb.epoch)
   122  						assert.NoError(t, err)
   123  					}
   124  				}()
   125  			}
   126  
   127  			consumerResults := make([][]int, test.consumers)
   128  			consumerGroup := &sync.WaitGroup{}
   129  			consumerGroup.Add(test.consumers)
   130  			for i := 0; i < test.consumers; i++ {
   131  				results := make([]int, test.itemsPerProducer)
   132  				consumerResults[i] = results
   133  				go func() {
   134  					defer consumerGroup.Done()
   135  					for {
   136  						item, err := rb.Pop()
   137  
   138  						if err != nil {
   139  							assert.Equal(t, io.EOF, err)
   140  							return
   141  						}
   142  
   143  						results[item.(int)]++
   144  					}
   145  				}()
   146  			}
   147  
   148  			producerGroup.Wait()
   149  			err := rb.Close()
   150  			assert.NoError(t, err)
   151  			consumerGroup.Wait()
   152  
   153  			for i := 0; i < test.itemsPerProducer; i++ {
   154  				sum := 0
   155  				for j := 0; j < test.consumers; j++ {
   156  					sum += consumerResults[j][i]
   157  				}
   158  
   159  				assert.Equal(t, test.producers, sum)
   160  			}
   161  		})
   162  	}
   163  }
   164  
   165  func TestRingBufferEpoch(t *testing.T) {
   166  	rb := NewRingBuffer(1024)
   167  	epoch := rb.Reset()
   168  	err := rb.Push(1, epoch)
   169  	assert.NoError(t, err)
   170  	err = rb.Push(2, epoch+1)
   171  	assert.Error(t, err)
   172  	assert.Equal(t, ErrWrongEpoch, err)
   173  	v, ok := rb.TryPop()
   174  	assert.True(t, ok)
   175  	assert.Equal(t, 1, v)
   176  	_, ok = rb.TryPop()
   177  	assert.False(t, ok)
   178  	newEpoch := rb.Reset()
   179  	assert.NotEqual(t, epoch, newEpoch)
   180  }