first commit

blocklock.go

@@ -0,0 +1,75 @@
+package expirelock
+
+import (
+	"context"
+	"errors"
+	"expirelock/os"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type Block struct {
+	token atomic.Value
+	mu    sync.Mutex
+	dmu   sync.Mutex
+	celFu context.CancelFunc
+}
+
+func NewBlock() *Block {
+	return &Block{}
+}
+
+func (b *Block) Lock(ttl time.Duration) error {
+	b.mu.Lock()
+
+	b.dmu.Lock()
+	defer b.dmu.Unlock()
+
+	token := os.GetCurrentProcessAndGogroutineIDStr()
+	b.token.Store(token)
+	b.celFu = nil
+
+	if ttl <= 0 {
+		return nil
+	}
+
+	var ctx context.Context
+	ctx, b.celFu = context.WithCancel(context.Background())
+	go func() {
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(ttl):
+			_ = b.unlock(token)
+		}
+	}()
+
+	return nil
+}
+
+func (b *Block) unlock(token string) error {
+	b.dmu.Lock()
+	defer b.dmu.Unlock()
+
+	id, _ := b.token.Load().(string)
+	if id != token {
+		return errors.New("invalid token")
+	}
+
+	if b.celFu != nil {
+		b.celFu()
+	}
+
+	b.token.Store("")
+
+	b.mu.Unlock()
+
+	return nil
+}
+
+func (b *Block) Unlock() error {
+	token := os.GetCurrentProcessAndGogroutineIDStr()
+
+	return b.unlock(token)
+}

go.mod

@@ -0,0 +1,3 @@
+module git.zhangshuocauc.cn/redhat/goexpirelock
+
+go 1.24.2

lock_test.go

@@ -0,0 +1,41 @@
+package expirelock
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestLock_Lock(t *testing.T) {
+	lock(NewBlock())
+}
+
+func TestLock_SpinLock(t *testing.T) {
+	lock(NewSpin())
+}
+
+func lock(lock ExpireLocker) {
+	wg := sync.WaitGroup{}
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func(i int) {
+			lock.Lock(1 * time.Second)
+			fmt.Printf("go: %d, lock\n", i)
+			time.Sleep(1 * time.Second)
+			fmt.Println(lock.Unlock())
+			fmt.Printf("go: %d, unlock\n", i)
+			wg.Done()
+		}(i)
+	}
+
+	lock.Lock(0)
+	fmt.Println(lock.Unlock())
+
+	//lock.Lock(1 * time.Second)
+	time.Sleep(1500 * time.Millisecond)
+	lock.Lock(1 * time.Second)
+	fmt.Println(lock.Unlock())
+
+	wg.Wait()
+}

locker.go

@@ -0,0 +1,8 @@
+package expirelock
+
+import "time"
+
+type ExpireLocker interface {
+	Lock(time.Duration) error
+	Unlock() error
+}

os/os.go

@@ -0,0 +1,28 @@
+package os
+
+import (
+	"fmt"
+	"os"
+	"runtime"
+	"strings"
+)
+
+// 获取由进程id+协程id组成的二位标识字符串
+func GetCurrentProcessAndGogroutineIDStr() string {
+	pid := GetCurrentProcessID()
+	goroutineID := GetCurrentGoroutineID()
+	return fmt.Sprintf("%d_%s", pid, goroutineID)
+}
+
+// 获取当前的协程id
+func GetCurrentGoroutineID() string {
+	buf := make([]byte, 128)
+	buf = buf[:runtime.Stack(buf, false)]
+	stackInfo := string(buf)
+	return strings.TrimSpace(strings.Split(strings.Split(stackInfo, "[running]")[0], "goroutine")[1])
+}
+
+// 获取当前的进程id
+func GetCurrentProcessID() int {
+	return os.Getpid()
+}

spinlock.go

@@ -0,0 +1,80 @@
+package expirelock
+
+import (
+	"errors"
+	"expirelock/os"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type Spin struct {
+	token  atomic.Value
+	mu     sync.Mutex
+	expire time.Time
+	locked bool
+}
+
+func NewSpin() *Spin {
+	return &Spin{}
+}
+
+func (s *Spin) tryLock(ttl time.Duration) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	now := time.Now()
+	if !s.locked || now.After(s.expire) {
+		s.locked = true
+		s.token.Store(os.GetCurrentProcessAndGogroutineIDStr())
+		s.expire = now.Add(ttl)
+
+		return nil
+	}
+
+	return errors.New("already locked")
+}
+
+func (s *Spin) Unlock() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	id, _ := s.token.Load().(string)
+	if id != os.GetCurrentProcessAndGogroutineIDStr() {
+		return errors.New("lock does not belong to this process")
+	}
+
+	if !s.locked || time.Now().After(s.expire) {
+		return errors.New("unlock already locked")
+	}
+
+	s.locked = false
+
+	return nil
+}
+
+func (s *Spin) Lock(ttl time.Duration) error {
+	err := s.tryLock(ttl)
+
+	if err == nil {
+		return nil
+	}
+
+	if ttl <= 0 {
+		return err
+	}
+
+	ticker := time.NewTicker(20 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-time.After(ttl):
+			return errors.New("lock timeout")
+		case <-ticker.C:
+			if err := s.tryLock(ttl); err == nil {
+				return nil
+			}
+		}
+	}
+}