github.com/xhghs/rclone@v1.51.1-0.20200430155106-e186a28cced8/lib/pool/pool_test.go (about)

     1  package pool
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"math/rand"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  // makes the allocations be unreliable
    14  func makeUnreliable(bp *Pool) {
    15  	bp.alloc = func(size int) ([]byte, error) {
    16  		if rand.Intn(3) != 0 {
    17  			return nil, errors.New("failed to allocate memory")
    18  		}
    19  		return make([]byte, size), nil
    20  	}
    21  	bp.free = func(b []byte) error {
    22  		if rand.Intn(3) != 0 {
    23  			return errors.New("failed to free memory")
    24  		}
    25  		return nil
    26  	}
    27  }
    28  
    29  func testGetPut(t *testing.T, useMmap bool, unreliable bool) {
    30  	bp := New(60*time.Second, 4096, 2, useMmap)
    31  	if unreliable {
    32  		makeUnreliable(bp)
    33  	}
    34  
    35  	assert.Equal(t, 0, bp.InUse())
    36  
    37  	b1 := bp.Get()
    38  	assert.Equal(t, 1, bp.InUse())
    39  	assert.Equal(t, 0, bp.InPool())
    40  	assert.Equal(t, 1, bp.Alloced())
    41  
    42  	b2 := bp.Get()
    43  	assert.Equal(t, 2, bp.InUse())
    44  	assert.Equal(t, 0, bp.InPool())
    45  	assert.Equal(t, 2, bp.Alloced())
    46  
    47  	b3 := bp.Get()
    48  	assert.Equal(t, 3, bp.InUse())
    49  	assert.Equal(t, 0, bp.InPool())
    50  	assert.Equal(t, 3, bp.Alloced())
    51  
    52  	bp.Put(b1)
    53  	assert.Equal(t, 2, bp.InUse())
    54  	assert.Equal(t, 1, bp.InPool())
    55  	assert.Equal(t, 3, bp.Alloced())
    56  
    57  	bp.Put(b2)
    58  	assert.Equal(t, 1, bp.InUse())
    59  	assert.Equal(t, 2, bp.InPool())
    60  	assert.Equal(t, 3, bp.Alloced())
    61  
    62  	bp.Put(b3)
    63  	assert.Equal(t, 0, bp.InUse())
    64  	assert.Equal(t, 2, bp.InPool())
    65  	assert.Equal(t, 2, bp.Alloced())
    66  
    67  	addr := func(b []byte) string {
    68  		return fmt.Sprintf("%p", &b[0])
    69  	}
    70  	b1a := bp.Get()
    71  	assert.Equal(t, addr(b2), addr(b1a))
    72  	assert.Equal(t, 1, bp.InUse())
    73  	assert.Equal(t, 1, bp.InPool())
    74  	assert.Equal(t, 2, bp.Alloced())
    75  
    76  	b2a := bp.Get()
    77  	assert.Equal(t, addr(b1), addr(b2a))
    78  	assert.Equal(t, 2, bp.InUse())
    79  	assert.Equal(t, 0, bp.InPool())
    80  	assert.Equal(t, 2, bp.Alloced())
    81  
    82  	bp.Put(b1a)
    83  	bp.Put(b2a)
    84  	assert.Equal(t, 0, bp.InUse())
    85  	assert.Equal(t, 2, bp.InPool())
    86  	assert.Equal(t, 2, bp.Alloced())
    87  
    88  	assert.Panics(t, func() {
    89  		bp.Put(make([]byte, 1))
    90  	})
    91  
    92  	bp.Flush()
    93  	assert.Equal(t, 0, bp.InUse())
    94  	assert.Equal(t, 0, bp.InPool())
    95  	assert.Equal(t, 0, bp.Alloced())
    96  }
    97  
    98  func testFlusher(t *testing.T, useMmap bool, unreliable bool) {
    99  	bp := New(50*time.Millisecond, 4096, 2, useMmap)
   100  	if unreliable {
   101  		makeUnreliable(bp)
   102  	}
   103  
   104  	b1 := bp.Get()
   105  	b2 := bp.Get()
   106  	b3 := bp.Get()
   107  	bp.Put(b1)
   108  	bp.Put(b2)
   109  	bp.Put(b3)
   110  	assert.Equal(t, 0, bp.InUse())
   111  	assert.Equal(t, 2, bp.InPool())
   112  	assert.Equal(t, 2, bp.Alloced())
   113  	bp.mu.Lock()
   114  	assert.Equal(t, 0, bp.minFill)
   115  	assert.Equal(t, true, bp.flushPending)
   116  	bp.mu.Unlock()
   117  
   118  	checkFlushHasHappened := func(desired int) {
   119  		var n int
   120  		for i := 0; i < 10; i++ {
   121  			time.Sleep(100 * time.Millisecond)
   122  			n = bp.InPool()
   123  			if n <= desired {
   124  				break
   125  			}
   126  		}
   127  		assert.Equal(t, desired, n)
   128  	}
   129  
   130  	checkFlushHasHappened(0)
   131  	assert.Equal(t, 0, bp.InUse())
   132  	assert.Equal(t, 0, bp.InPool())
   133  	assert.Equal(t, 0, bp.Alloced())
   134  	bp.mu.Lock()
   135  	assert.Equal(t, 0, bp.minFill)
   136  	assert.Equal(t, false, bp.flushPending)
   137  	bp.mu.Unlock()
   138  
   139  	// Now do manual aging to check it is working properly
   140  	bp = New(100*time.Second, 4096, 2, useMmap)
   141  
   142  	// Check the new one doesn't get flushed
   143  	b1 = bp.Get()
   144  	b2 = bp.Get()
   145  	bp.Put(b1)
   146  	bp.Put(b2)
   147  
   148  	bp.mu.Lock()
   149  	assert.Equal(t, 0, bp.minFill)
   150  	assert.Equal(t, true, bp.flushPending)
   151  	bp.mu.Unlock()
   152  
   153  	bp.flushAged()
   154  
   155  	assert.Equal(t, 0, bp.InUse())
   156  	assert.Equal(t, 2, bp.InPool())
   157  	assert.Equal(t, 2, bp.Alloced())
   158  	bp.mu.Lock()
   159  	assert.Equal(t, 2, bp.minFill)
   160  	assert.Equal(t, true, bp.flushPending)
   161  	bp.mu.Unlock()
   162  
   163  	bp.Put(bp.Get())
   164  
   165  	assert.Equal(t, 0, bp.InUse())
   166  	assert.Equal(t, 2, bp.InPool())
   167  	assert.Equal(t, 2, bp.Alloced())
   168  	bp.mu.Lock()
   169  	assert.Equal(t, 1, bp.minFill)
   170  	assert.Equal(t, true, bp.flushPending)
   171  	bp.mu.Unlock()
   172  
   173  	bp.flushAged()
   174  
   175  	assert.Equal(t, 0, bp.InUse())
   176  	assert.Equal(t, 1, bp.InPool())
   177  	assert.Equal(t, 1, bp.Alloced())
   178  	bp.mu.Lock()
   179  	assert.Equal(t, 1, bp.minFill)
   180  	assert.Equal(t, true, bp.flushPending)
   181  	bp.mu.Unlock()
   182  
   183  	bp.flushAged()
   184  
   185  	assert.Equal(t, 0, bp.InUse())
   186  	assert.Equal(t, 0, bp.InPool())
   187  	assert.Equal(t, 0, bp.Alloced())
   188  	bp.mu.Lock()
   189  	assert.Equal(t, 0, bp.minFill)
   190  	assert.Equal(t, false, bp.flushPending)
   191  	bp.mu.Unlock()
   192  }
   193  
   194  func TestPool(t *testing.T) {
   195  	for _, test := range []struct {
   196  		name       string
   197  		useMmap    bool
   198  		unreliable bool
   199  	}{
   200  		{
   201  			name:       "make",
   202  			useMmap:    false,
   203  			unreliable: false,
   204  		},
   205  		{
   206  			name:       "mmap",
   207  			useMmap:    true,
   208  			unreliable: false,
   209  		},
   210  		{
   211  			name:       "canFail",
   212  			useMmap:    false,
   213  			unreliable: true,
   214  		},
   215  	} {
   216  		t.Run(test.name, func(t *testing.T) {
   217  			t.Run("GetPut", func(t *testing.T) { testGetPut(t, test.useMmap, test.unreliable) })
   218  			t.Run("Flusher", func(t *testing.T) { testFlusher(t, test.useMmap, test.unreliable) })
   219  		})
   220  	}
   221  }