github.com/sequix/cortex@v1.1.6/pkg/chunk/cache/memcached_test.go (about)

     1  package cache_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"sync/atomic"
     7  	"testing"
     8  
     9  	"github.com/bradfitz/gomemcache/memcache"
    10  	"github.com/sequix/cortex/pkg/chunk/cache"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestMemcached(t *testing.T) {
    15  	t.Run("unbatched", func(t *testing.T) {
    16  		client := newMockMemcache()
    17  		memcache := cache.NewMemcached(cache.MemcachedConfig{}, client, "test")
    18  
    19  		testMemcache(t, memcache)
    20  	})
    21  
    22  	t.Run("batched", func(t *testing.T) {
    23  		client := newMockMemcache()
    24  		memcache := cache.NewMemcached(cache.MemcachedConfig{
    25  			BatchSize:   10,
    26  			Parallelism: 5,
    27  		}, client, "test")
    28  
    29  		testMemcache(t, memcache)
    30  	})
    31  }
    32  
    33  func testMemcache(t *testing.T, memcache *cache.Memcached) {
    34  	numKeys := 1000
    35  
    36  	ctx := context.Background()
    37  	keysIncMissing := make([]string, 0, numKeys)
    38  	keys := make([]string, 0, numKeys)
    39  	bufs := make([][]byte, 0, numKeys)
    40  
    41  	// Insert 1000 keys skipping all multiples of 5.
    42  	for i := 0; i < numKeys; i++ {
    43  		keysIncMissing = append(keysIncMissing, string(i))
    44  		if i%5 == 0 {
    45  			continue
    46  		}
    47  
    48  		keys = append(keys, string(i))
    49  		bufs = append(bufs, []byte(string(i)))
    50  	}
    51  	memcache.Store(ctx, keys, bufs)
    52  
    53  	found, bufs, missing := memcache.Fetch(ctx, keysIncMissing)
    54  	for i := 0; i < numKeys; i++ {
    55  		if i%5 == 0 {
    56  			require.Equal(t, string(i), missing[0])
    57  			missing = missing[1:]
    58  			continue
    59  		}
    60  
    61  		require.Equal(t, string(i), found[0])
    62  		require.Equal(t, string(i), string(bufs[0]))
    63  		found = found[1:]
    64  		bufs = bufs[1:]
    65  	}
    66  }
    67  
    68  // mockMemcache whose calls fail 1/3rd of the time.
    69  type mockMemcacheFailing struct {
    70  	*mockMemcache
    71  	calls uint64
    72  }
    73  
    74  func newMockMemcacheFailing() *mockMemcacheFailing {
    75  	return &mockMemcacheFailing{
    76  		mockMemcache: newMockMemcache(),
    77  	}
    78  }
    79  
    80  func (c *mockMemcacheFailing) GetMulti(keys []string) (map[string]*memcache.Item, error) {
    81  	calls := atomic.AddUint64(&c.calls, 1)
    82  	if calls%3 == 0 {
    83  		return nil, errors.New("fail")
    84  	}
    85  
    86  	return c.mockMemcache.GetMulti(keys)
    87  }
    88  
    89  func TestMemcacheFailure(t *testing.T) {
    90  	t.Run("unbatched", func(t *testing.T) {
    91  		client := newMockMemcacheFailing()
    92  		memcache := cache.NewMemcached(cache.MemcachedConfig{}, client, "test")
    93  
    94  		testMemcacheFailing(t, memcache)
    95  	})
    96  
    97  	t.Run("batched", func(t *testing.T) {
    98  		client := newMockMemcacheFailing()
    99  		memcache := cache.NewMemcached(cache.MemcachedConfig{
   100  			BatchSize:   10,
   101  			Parallelism: 5,
   102  		}, client, "test")
   103  
   104  		testMemcacheFailing(t, memcache)
   105  	})
   106  }
   107  
   108  func testMemcacheFailing(t *testing.T, memcache *cache.Memcached) {
   109  	numKeys := 1000
   110  
   111  	ctx := context.Background()
   112  	keysIncMissing := make([]string, 0, numKeys)
   113  	keys := make([]string, 0, numKeys)
   114  	bufs := make([][]byte, 0, numKeys)
   115  	// Insert 1000 keys skipping all multiples of 5.
   116  	for i := 0; i < numKeys; i++ {
   117  		keysIncMissing = append(keysIncMissing, string(i))
   118  		if i%5 == 0 {
   119  			continue
   120  		}
   121  		keys = append(keys, string(i))
   122  		bufs = append(bufs, []byte(string(i)))
   123  	}
   124  	memcache.Store(ctx, keys, bufs)
   125  
   126  	for i := 0; i < 10; i++ {
   127  		found, bufs, missing := memcache.Fetch(ctx, keysIncMissing)
   128  
   129  		require.Equal(t, len(found), len(bufs))
   130  		for i := range found {
   131  			require.Equal(t, found[i], string(bufs[i]))
   132  		}
   133  
   134  		keysReturned := make(map[string]struct{})
   135  		for _, key := range found {
   136  			_, ok := keysReturned[key]
   137  			require.False(t, ok, "duplicate key returned")
   138  
   139  			keysReturned[key] = struct{}{}
   140  		}
   141  		for _, key := range missing {
   142  			_, ok := keysReturned[key]
   143  			require.False(t, ok, "duplicate key returned")
   144  
   145  			keysReturned[key] = struct{}{}
   146  		}
   147  
   148  		for _, key := range keys {
   149  			_, ok := keysReturned[key]
   150  			require.True(t, ok, "key missing %s", key)
   151  		}
   152  	}
   153  }