New interface TryLockWithContext and RTryLockWithContext

cas.go

@@ -22,6 +22,9 @@ type CASMutex interface {
 	// TryLockWithTimeout attempts to acquire the write lock within a period of time.
 	// Return false if spending time is more than duration and no chance to acquire it.
 	TryLockWithTimeout(time.Duration) bool
+	// TryLockWithContext attempts to acquire the write lock, blocking until resources
+	// are available or ctx is done (timeout or cancellation)
+	TryLockWithContext(ctx context.Context) bool
 	// Unlock releases the write lock
 	Unlock()
 
@@ -34,6 +37,9 @@ type CASMutex interface {
 	// RTryLockWithTimeout attempts to acquire the read lock within a period of time.
 	// Return false if spending time is more than duration and no chance to acquire it.
 	RTryLockWithTimeout(time.Duration) bool
+	// RTryLockWithContext attempts to acquire the read lock, blocking until resources
+	// are available or ctx is done (timeout or cancellation)
+	RTryLockWithContext(ctx context.Context) bool
 	// RUnlock releases the read lock
 	RUnlock()
 }
@@ -87,14 +93,6 @@ func (m *casMutex) broadcast() {
 	close(ch)
 }
 
-func (m *casMutex) Lock() {
-	ctx := context.Background()
-	m.turnstile.Acquire(ctx, 1)
-	defer m.turnstile.Release(1)
-
-	m.tryLock(ctx)
-}
-
 func (m *casMutex) tryLock(ctx context.Context) bool {
 	for {
 		broker := m.listen()
@@ -112,14 +110,31 @@ func (m *casMutex) tryLock(ctx context.Context) bool {
 
 		select {
 		case <-ctx.Done():
-			// timeout
+			// timeout or cancellation
 			return false
 		case <-broker:
 			// waiting for signal triggered by m.broadcast() and trying again.
 		}
 	}
 }
 
+func (m *casMutex) TryLockWithContext(ctx context.Context) bool {
+	if err := m.turnstile.Acquire(ctx, 1); err != nil {
+		// Acquire failed due to timeout or cancellation
+		return false
+	}
+
+	defer m.turnstile.Release(1)
+
+	return m.tryLock(ctx)
+}
+
+func (m *casMutex) Lock() {
+	ctx := context.Background()
+
+	m.TryLockWithContext(ctx)
+}
+
 func (m *casMutex) TryLock() bool {
 	if !m.turnstile.TryAcquire(1) {
 		return false
@@ -134,14 +149,7 @@ func (m *casMutex) TryLockWithTimeout(duration time.Duration) bool {
 	ctx, cancel := context.WithTimeout(context.Background(), duration)
 	defer cancel()
 
-	if err := m.turnstile.Acquire(ctx, 1); err != nil {
-		// Acquire failed due to timeout
-		return false
-	}
-
-	defer m.turnstile.Release(1)
-
-	return m.tryLock(ctx)
+	return m.TryLockWithContext(ctx)
 }
 
 func (m *casMutex) Unlock() {
@@ -156,14 +164,6 @@ func (m *casMutex) Unlock() {
 	m.broadcast()
 }
 
-func (m *casMutex) RLock() {
-	ctx := context.Background()
-	m.turnstile.Acquire(ctx, 1)
-	m.turnstile.Release(1)
-
-	m.rTryLock(ctx)
-}
-
 func (m *casMutex) rTryLock(ctx context.Context) bool {
 	for {
 		broker := m.listen()
@@ -182,7 +182,7 @@ func (m *casMutex) rTryLock(ctx context.Context) bool {
 
 		select {
 		case <-ctx.Done():
-			// timeout
+			// timeout or cancellation
 			return false
 		default:
 			switch st {
@@ -195,14 +195,31 @@ func (m *casMutex) rTryLock(ctx context.Context) bool {
 
 		select {
 		case <-ctx.Done():
-			// timeout
+			// timeout or cancellation
 			return false
 		case <-broker:
 			// waiting for signal triggered by m.broadcast() and trying again.
 		}
 	}
 }
 
+func (m *casMutex) RTryLockWithContext(ctx context.Context) bool {
+	if err := m.turnstile.Acquire(ctx, 1); err != nil {
+		// Acquire failed due to timeout or cancellation
+		return false
+	}
+
+	m.turnstile.Release(1)
+
+	return m.rTryLock(ctx)
+}
+
+func (m *casMutex) RLock() {
+	ctx := context.Background()
+
+	m.RTryLockWithContext(ctx)
+}
+
 func (m *casMutex) RTryLock() bool {
 	if !m.turnstile.TryAcquire(1) {
 		return false
@@ -217,14 +234,7 @@ func (m *casMutex) RTryLockWithTimeout(duration time.Duration) bool {
 	ctx, cancel := context.WithTimeout(context.Background(), duration)
 	defer cancel()
 
-	if err := m.turnstile.Acquire(ctx, 1); err != nil {
-		// Acquire failed due to timeout
-		return false
-	}
-
-	m.turnstile.Release(1)
-
-	return m.rTryLock(ctx)
+	return m.RTryLockWithContext(ctx)
 }
 
 func (m *casMutex) RUnlock() {

chan.go

@@ -1,6 +1,9 @@
 package lock
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 // ChanMutex provides interfaces of spinlock and trylock implemented by channel.
 type ChanMutex interface {
@@ -13,6 +16,9 @@ type ChanMutex interface {
 	// TryLockWithTimeout attempts to acquire the lock within a period of time.
 	// Return false if spending time is more than duration and no chance to acquire it.
 	TryLockWithTimeout(time.Duration) bool
+	// TryLockWithContext attempts to acquire the lock, blocking until resources
+	// are available or ctx is done (timeout or cancellation)
+	TryLockWithContext(ctx context.Context) bool
 	// Unlock releases the lock
 	Unlock()
 }
@@ -38,15 +44,23 @@ func (m *chanMutex) TryLock() bool {
 	}
 }
 
-func (m *chanMutex) TryLockWithTimeout(duration time.Duration) bool {
+func (m *chanMutex) TryLockWithContext(ctx context.Context) bool {
 	select {
 	case m.lockChan <- struct{}{}:
 		return true
-	case <-time.After(duration):
+	case <-ctx.Done():
+		// timeout or cancellation
 		return false
 	}
 }
 
+func (m *chanMutex) TryLockWithTimeout(duration time.Duration) bool {
+	ctx, cancel := context.WithTimeout(context.Background(), duration)
+	defer cancel()
+
+	return m.TryLockWithContext(ctx)
+}
+
 // NewChanMutex returns ChanMutex lock
 func NewChanMutex() ChanMutex {
 	return &chanMutex{