Self expirying lock in golang

timedlock.go

package main

import (
	"fmt"
	"sync"
	"time"
)

type TimedLock struct {
	mutex chan (uint)
	timer *time.Timer
}

func NewTimedLock() *TimedLock {
	l := &TimedLock{
		mutex: make(chan (uint), 1),
	}
	l.mutex <- 1

	return l
}

func (l *TimedLock) LockFor(ttl time.Duration, abort <-chan (uint)) {
	select {
	case <-abort:
	case <-l.mutex:
		l.timer = time.AfterFunc(ttl, func() {
			l.mutex <- 1
		})
	}
}

func (l *TimedLock) Unlock() {
	l.timer.Stop()
	l.mutex <- 1
}

type LockMap struct {
	sync.RWMutex

	locks map[string]*TimedLock
}

func NewLockMap() *LockMap {
	return &LockMap{
		locks: make(map[string]*TimedLock),
	}
}

func (lm *LockMap) AcquireLock(key string, ttl time.Duration, timeout time.Duration) bool {
	lm.Lock()
	if _, ok := lm.locks[key]; !ok {
		lm.locks[key] = NewTimedLock()
	}
	lm.Unlock()

	abort := make(chan uint)
	acquired := make(chan uint)

	go func(key string, ttl time.Duration) {
		lm.locks[key].LockFor(ttl, abort)
		close(acquired)
	}(key, ttl)

	t := time.NewTimer(timeout)
	select {
	case <-t.C:
		close(abort)
		return false // timeout reached
	case <-acquired:
		return true
	}
}

func (lm *LockMap) ReleaseLock(key string) {
	lm.RLock()
	defer lm.RUnlock()

	lm.locks[key].Unlock()
}

func main() {

	lm := NewLockMap()

	var wg sync.WaitGroup

	fmt.Println(time.Now(), "ready, set, go!")

	// Lock should success immediately
	go func(k string, id int) {
		if ok := lm.AcquireLock(k, time.Second*3, time.Second*3); ok {
			fmt.Println(time.Now(), "goroutine", id, "acquired", k)
		} else {
			fmt.Println(time.Now(), "goroutine", id, "cannot acquired", k)
		}

		<-time.NewTimer(time.Second * 1).C

		lm.ReleaseLock(k)
		fmt.Println(time.Now(), "goroutine", id, "released", k)

		wg.Done()
	}("KEY", 1)
	wg.Add(1)

	// Delay starting next goroutine
	<-time.NewTimer(time.Millisecond * 10).C

	// Lock have to wait for 1s
	go func(k string, id int) {
		if ok := lm.AcquireLock(k, time.Second*3, time.Second*3); ok {
			fmt.Println(time.Now(), "goroutine", id, "acquired", k)
		} else {
			fmt.Println(time.Now(), "goroutine", id, "cannot acquired", k)
		}

		wg.Done()
	}("KEY", 2)
	wg.Add(1)

	// Delay starting next goroutine
	<-time.NewTimer(time.Millisecond * 10).C

	// Lock should failed at 3s as goroutine 1 took 1s and goroutine 2 took 3s
	go func(k string, id int) {
		if ok := lm.AcquireLock(k, time.Second*3, time.Second*3); ok {
			fmt.Println(time.Now(), "goroutine", id, "acquired", k)
		} else {
			fmt.Println(time.Now(), "goroutine", id, "cannot acquired", k)
		}

		wg.Done()
	}("KEY", 3)
	wg.Add(1)

	// Delay starting next goroutine
	<-time.NewTimer(time.Millisecond * 10).C

	// Lock should success at 4s as goroutine 1 took 1s and goroutine 2 took 3s
	go func(k string, id int) {
		if ok := lm.AcquireLock(k, time.Second*3, time.Second*4); ok {
			fmt.Println(time.Now(), "goroutine", id, "acquired", k)
		} else {
			fmt.Println(time.Now(), "goroutine", id, "cannot acquired", k)
		}

		wg.Done()
	}("KEY", 4)
	wg.Add(1)

	wg.Wait()
}