github.com/keybase/client/go@v0.0.0-20240309051027-028f7c731f8b/kbfs/libkbfs/semaphore_disk_limiter_test.go (about)

     1  // Copyright 2017 Keybase Inc. All rights reserved.
     2  // Use of this source code is governed by a BSD
     3  // license that can be found in the LICENSE file.
     4  
     5  package libkbfs
     6  
     7  import (
     8  	"context"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/keybase/client/go/protocol/keybase1"
    13  	"github.com/pkg/errors"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  func TestSemaphoreDiskLimiterBlockBasic(t *testing.T) {
    18  	sdl := newSemaphoreDiskLimiter(10, 2, 12)
    19  
    20  	ctx := context.Background()
    21  
    22  	u := keybase1.UserOrTeamID("")
    23  	availBytes, availFiles, err := sdl.reserveWithBackpressure(ctx,
    24  		journalLimitTrackerType, 9, 1, u)
    25  	require.NoError(t, err)
    26  	require.Equal(t, int64(1), availBytes)
    27  	require.Equal(t, int64(1), availFiles)
    28  
    29  	require.Equal(t, int64(1), sdl.byteSemaphore.Count())
    30  	require.Equal(t, int64(1), sdl.fileSemaphore.Count())
    31  
    32  	usedQuotaBytes, quotaBytes := sdl.getQuotaInfo(u)
    33  	require.Equal(t, int64(0), usedQuotaBytes)
    34  	require.Equal(t, int64(12), quotaBytes)
    35  
    36  	sdl.commitOrRollback(ctx, journalLimitTrackerType, 9, 1, true, u)
    37  
    38  	require.Equal(t, int64(1), sdl.byteSemaphore.Count())
    39  	require.Equal(t, int64(1), sdl.fileSemaphore.Count())
    40  
    41  	usedQuotaBytes, quotaBytes = sdl.getQuotaInfo(u)
    42  	require.Equal(t, int64(9), usedQuotaBytes)
    43  	require.Equal(t, int64(12), quotaBytes)
    44  
    45  	sdl.onBlocksFlush(ctx, 9, u)
    46  
    47  	require.Equal(t, int64(1), sdl.byteSemaphore.Count())
    48  	require.Equal(t, int64(1), sdl.fileSemaphore.Count())
    49  
    50  	usedQuotaBytes, quotaBytes = sdl.getQuotaInfo(u)
    51  	require.Equal(t, int64(0), usedQuotaBytes)
    52  	require.Equal(t, int64(12), quotaBytes)
    53  
    54  	sdl.release(ctx, unknownLimitTrackerType, 9, 1)
    55  
    56  	require.Equal(t, int64(10), sdl.byteSemaphore.Count())
    57  	require.Equal(t, int64(2), sdl.fileSemaphore.Count())
    58  
    59  	usedQuotaBytes, quotaBytes = sdl.getQuotaInfo(u)
    60  	require.Equal(t, int64(0), usedQuotaBytes)
    61  	require.Equal(t, int64(12), quotaBytes)
    62  }
    63  
    64  // TestSemaphoreDiskLimiterBeforeBlockPutError checks that
    65  // semaphoreDiskLimiter.beforeBlockPut handles errors correctly; in
    66  // particular, that we don't leak either bytes or files if either
    67  // semaphore times out.
    68  func TestSemaphoreDiskLimiterBeforeBlockPutError(t *testing.T) {
    69  	sdl := newSemaphoreDiskLimiter(10, 1, 12)
    70  
    71  	ctx, cancel := context.WithTimeout(
    72  		context.Background(), 3*time.Millisecond)
    73  	defer cancel()
    74  
    75  	availBytes, availFiles, err := sdl.reserveWithBackpressure(ctx, journalLimitTrackerType, 10, 2,
    76  		keybase1.UserOrTeamID(""))
    77  	require.Equal(t, context.DeadlineExceeded, errors.Cause(err))
    78  	require.Equal(t, int64(10), availBytes)
    79  	require.Equal(t, int64(1), availFiles)
    80  
    81  	require.Equal(t, int64(10), sdl.byteSemaphore.Count())
    82  	require.Equal(t, int64(1), sdl.fileSemaphore.Count())
    83  }