github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/flock/flock_test.go (about)

     1  package flock_test
     2  
     3  import (
     4  	"context"
     5  	"io/ioutil"
     6  	"sync/atomic"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/Schaudge/grailbase/flock"
    11  	"github.com/grailbio/testutil/assert"
    12  )
    13  
    14  func TestLock(t *testing.T) {
    15  	tempDir, err := ioutil.TempDir("", "")
    16  	if err != nil {
    17  		t.Fatal(err)
    18  	}
    19  
    20  	lockPath := tempDir + "/lock"
    21  	lock := flock.New(lockPath)
    22  
    23  	// Test uncontended locks
    24  	ctx := context.Background()
    25  	for i := 0; i < 3; i++ {
    26  		assert.NoError(t, lock.Lock(ctx))
    27  		assert.NoError(t, lock.Unlock())
    28  	}
    29  
    30  	assert.NoError(t, lock.Lock(ctx))
    31  	locked := int64(0)
    32  	doneCh := make(chan struct{})
    33  	go func() {
    34  		assert.NoError(t, lock.Lock(ctx))
    35  		atomic.StoreInt64(&locked, 1)
    36  		assert.NoError(t, lock.Unlock())
    37  		atomic.StoreInt64(&locked, 2)
    38  		doneCh <- struct{}{}
    39  	}()
    40  
    41  	time.Sleep(500 * time.Millisecond)
    42  	if atomic.LoadInt64(&locked) != 0 {
    43  		t.Errorf("locked=%d", locked)
    44  	}
    45  
    46  	assert.NoError(t, lock.Unlock())
    47  	<-doneCh
    48  	if atomic.LoadInt64(&locked) != 2 {
    49  		t.Errorf("locked=%d", locked)
    50  	}
    51  }
    52  
    53  func TestLockContext(t *testing.T) {
    54  	tempDir, err := ioutil.TempDir("", "")
    55  	if err != nil {
    56  		t.Fatal(err)
    57  	}
    58  	lockPath := tempDir + "/lock"
    59  
    60  	lock := flock.New(lockPath)
    61  	ctx := context.Background()
    62  	ctx2, cancel2 := context.WithCancel(ctx)
    63  	assert.NoError(t, lock.Lock(ctx2))
    64  	assert.NoError(t, lock.Unlock())
    65  
    66  	assert.NoError(t, lock.Lock(ctx))
    67  	go func() {
    68  		time.Sleep(500 * time.Millisecond)
    69  		cancel2()
    70  	}()
    71  	assert.Regexp(t, lock.Lock(ctx2), "context canceled")
    72  
    73  	assert.NoError(t, lock.Unlock())
    74  	// Make sure the lock is in a sane state by cycling lock-unlock again.
    75  	assert.NoError(t, lock.Lock(ctx))
    76  	assert.NoError(t, lock.Unlock())
    77  }