vitess.io/vitess@v0.16.2/go/vt/vtorc/logic/orchestrator_test.go (about)

     1  package logic
     2  
     3  import (
     4  	"sync/atomic"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestWaitForLocksRelease(t *testing.T) {
    12  	oldShutdownWaitTime := shutdownWaitTime
    13  	// Restore initial values
    14  	defer func() {
    15  		shutdownWaitTime = oldShutdownWaitTime
    16  	}()
    17  
    18  	t.Run("No locks to wait for", func(t *testing.T) {
    19  		// Initially when shardsLockCounter is zero, waitForLocksRelease should run immediately
    20  		timeSpent := waitForLocksReleaseAndGetTimeWaitedFor()
    21  		assert.Less(t, timeSpent, 1*time.Second, "waitForLocksRelease should run immediately if there are no locks to wait for")
    22  	})
    23  
    24  	t.Run("Timeout from shutdownWaitTime", func(t *testing.T) {
    25  		// Increment shardsLockCounter to simulate locking of a shard
    26  		atomic.AddInt32(&shardsLockCounter, +1)
    27  		defer func() {
    28  			// Restore the initial value
    29  			atomic.StoreInt32(&shardsLockCounter, 0)
    30  		}()
    31  		shutdownWaitTime = 200 * time.Millisecond
    32  		timeSpent := waitForLocksReleaseAndGetTimeWaitedFor()
    33  		assert.Greater(t, timeSpent, 100*time.Millisecond, "waitForLocksRelease should timeout after 200 milliseconds and not before")
    34  		assert.Less(t, timeSpent, 300*time.Millisecond, "waitForLocksRelease should timeout after 200 milliseconds and not take any longer")
    35  	})
    36  
    37  	t.Run("Successful wait for locks release", func(t *testing.T) {
    38  		// Increment shardsLockCounter to simulate locking of a shard
    39  		atomic.AddInt32(&shardsLockCounter, +1)
    40  		shutdownWaitTime = 500 * time.Millisecond
    41  		// Release the locks after 200 milliseconds
    42  		go func() {
    43  			time.Sleep(200 * time.Millisecond)
    44  			atomic.StoreInt32(&shardsLockCounter, 0)
    45  		}()
    46  		timeSpent := waitForLocksReleaseAndGetTimeWaitedFor()
    47  		assert.Greater(t, timeSpent, 100*time.Millisecond, "waitForLocksRelease should wait for the locks and not return early")
    48  		assert.Less(t, timeSpent, 300*time.Millisecond, "waitForLocksRelease should be successful after 200 milliseconds as all the locks are released")
    49  	})
    50  }
    51  
    52  func waitForLocksReleaseAndGetTimeWaitedFor() time.Duration {
    53  	start := time.Now()
    54  	waitForLocksRelease()
    55  	return time.Since(start)
    56  }