// Package keyedlock provides one-slot semaphores by key, shared by the
// providers that serialise work per instance name or per cached image
// (kvm2, proxmox).
package keyedlock

import (
	"context"
	"sync"
)

// Locks is a set of one-slot semaphores by key: a semaphore rather than a
// mutex so a waiter can give up with its ctx. Each entry counts the callers
// holding or waiting for it and is dropped from the map at zero, so the map
// does not grow with every key ever seen while two waiters on one key still
// always share one semaphore. The zero value is not usable; call New.
type Locks struct {
	mu    sync.Mutex
	locks map[string]*lock
}

type lock struct {
	sem  chan struct{}
	refs int
}

func New() *Locks {
	return &Locks{locks: make(map[string]*lock)}
}

// Acquire takes the lock for key, or gives up with ctx's error. The
// returned func releases it and must be called exactly once.
func (k *Locks) Acquire(ctx context.Context, key string) (release func(), err error) {
	k.mu.Lock()
	l, ok := k.locks[key]
	if !ok {
		l = &lock{sem: make(chan struct{}, 1)}
		k.locks[key] = l
	}
	l.refs++
	k.mu.Unlock()
	select {
	case l.sem <- struct{}{}:
		return func() {
			<-l.sem
			k.unref(key, l)
		}, nil
	case <-ctx.Done():
		k.unref(key, l)
		return nil, ctx.Err()
	}
}

func (k *Locks) unref(key string, l *lock) {
	k.mu.Lock()
	defer k.mu.Unlock()
	l.refs--
	if l.refs == 0 {
		delete(k.locks, key)
	}
}