github.com/outcaste-io/ristretto@v0.2.3/z/allocator_test.go (about)

     1  /*
     2   * Copyright 2020 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package z
    18  
    19  import (
    20  	"math/rand"
    21  	"sort"
    22  	"sync"
    23  	"testing"
    24  	"unsafe"
    25  
    26  	"github.com/stretchr/testify/require"
    27  )
    28  
    29  func TestAllocate(t *testing.T) {
    30  	a := NewAllocator(1024, "test")
    31  	defer a.Release()
    32  
    33  	check := func() {
    34  		t.Logf("Running checks\n")
    35  		require.Equal(t, 0, len(a.Allocate(0)))
    36  		require.Equal(t, 1, len(a.Allocate(1)))
    37  		require.Equal(t, 1<<20+1, len(a.Allocate(1<<20+1)))
    38  		require.Equal(t, 256<<20, len(a.Allocate(256<<20)))
    39  		require.Panics(t, func() { a.Allocate(maxAlloc + 1) })
    40  	}
    41  
    42  	check()
    43  	t.Logf("%s", a)
    44  	prev := a.Allocated()
    45  	t.Logf("Resetting\n")
    46  	a.Reset()
    47  	check()
    48  	t.Logf("%s", a)
    49  	require.Equal(t, int(prev), int(a.Allocated()))
    50  	t.Logf("Allocated: %d\n", prev)
    51  }
    52  
    53  func TestAllocateSize(t *testing.T) {
    54  	a := NewAllocator(1024, "test")
    55  	require.Equal(t, 1024, len(a.buffers[0]))
    56  	a.Release()
    57  
    58  	b := NewAllocator(1025, "test")
    59  	require.Equal(t, 2048, len(b.buffers[0]))
    60  	b.Release()
    61  }
    62  
    63  func TestAllocateReset(t *testing.T) {
    64  	a := NewAllocator(16, "test")
    65  	defer a.Release()
    66  
    67  	buf := make([]byte, 128)
    68  	rand.Read(buf)
    69  	for i := 0; i < 1000; i++ {
    70  		a.Copy(buf)
    71  	}
    72  
    73  	prev := a.Allocated()
    74  	a.Reset()
    75  	for i := 0; i < 100; i++ {
    76  		a.Copy(buf)
    77  	}
    78  	t.Logf("%s", a)
    79  	require.Equal(t, prev, a.Allocated())
    80  }
    81  
    82  func TestAllocateTrim(t *testing.T) {
    83  	a := NewAllocator(16, "test")
    84  	defer a.Release()
    85  
    86  	buf := make([]byte, 128)
    87  	rand.Read(buf)
    88  	for i := 0; i < 1000; i++ {
    89  		a.Copy(buf)
    90  	}
    91  
    92  	N := 2048
    93  	a.TrimTo(N)
    94  	require.LessOrEqual(t, int(a.Allocated()), N)
    95  }
    96  
    97  func TestPowTwo(t *testing.T) {
    98  	require.Equal(t, 2, log2(4))
    99  	require.Equal(t, 2, log2(7))
   100  	require.Equal(t, 3, log2(8))
   101  	require.Equal(t, 3, log2(15))
   102  	require.Equal(t, 4, log2(16))
   103  	require.Equal(t, 4, log2(31))
   104  	require.Equal(t, 10, log2(1024))
   105  	require.Equal(t, 10, log2(1025))
   106  	require.Equal(t, 10, log2(2047))
   107  	require.Equal(t, 11, log2(2048))
   108  }
   109  
   110  func TestAllocateAligned(t *testing.T) {
   111  	a := NewAllocator(1024, "test")
   112  	defer a.Release()
   113  
   114  	a.Allocate(1)
   115  	out := a.Allocate(1)
   116  	ptr := uintptr(unsafe.Pointer(&out[0]))
   117  	require.True(t, ptr%8 == 1)
   118  
   119  	out = a.AllocateAligned(5)
   120  	ptr = uintptr(unsafe.Pointer(&out[0]))
   121  	require.True(t, ptr%8 == 0)
   122  
   123  	out = a.AllocateAligned(3)
   124  	ptr = uintptr(unsafe.Pointer(&out[0]))
   125  	require.True(t, ptr%8 == 0)
   126  }
   127  
   128  func TestAllocateConcurrent(t *testing.T) {
   129  	a := NewAllocator(63, "test")
   130  	defer a.Release()
   131  
   132  	N := 10240
   133  	M := 16
   134  	var wg sync.WaitGroup
   135  
   136  	m := make(map[uintptr]struct{})
   137  	mu := new(sync.Mutex)
   138  	for i := 0; i < M; i++ {
   139  		wg.Add(1)
   140  		go func() {
   141  			defer wg.Done()
   142  			var bufs []uintptr
   143  			for j := 0; j < N; j++ {
   144  				buf := a.Allocate(16)
   145  				require.Equal(t, 16, len(buf))
   146  				bufs = append(bufs, uintptr(unsafe.Pointer(&buf[0])))
   147  			}
   148  
   149  			mu.Lock()
   150  			for _, b := range bufs {
   151  				if _, ok := m[b]; ok {
   152  					t.Fatalf("Did not expect to see the same ptr")
   153  				}
   154  				m[b] = struct{}{}
   155  			}
   156  			mu.Unlock()
   157  		}()
   158  	}
   159  	wg.Wait()
   160  	t.Logf("Size of allocator: %v. Allocator: %s\n", a.Size(), a)
   161  
   162  	require.Equal(t, N*M, len(m))
   163  	var sorted []uintptr
   164  	for ptr := range m {
   165  		sorted = append(sorted, ptr)
   166  	}
   167  
   168  	sort.Slice(sorted, func(i, j int) bool {
   169  		return sorted[i] < sorted[j]
   170  	})
   171  
   172  	var last uintptr
   173  	for _, ptr := range sorted {
   174  		if ptr-last < 16 {
   175  			t.Fatalf("Should not have less than 16: %v %v\n", ptr, last)
   176  		}
   177  		// fmt.Printf("ptr [%d]: %x %d\n", i, ptr, ptr-last)
   178  		last = ptr
   179  	}
   180  }
   181  
   182  func BenchmarkAllocate(b *testing.B) {
   183  	a := NewAllocator(15, "test")
   184  	b.RunParallel(func(pb *testing.PB) {
   185  		for pb.Next() {
   186  			buf := a.Allocate(1)
   187  			if len(buf) != 1 {
   188  				b.FailNow()
   189  			}
   190  		}
   191  	})
   192  	b.StopTimer()
   193  	b.Logf("%s", a)
   194  }