PaddedCell/lock-mgr/main.go

package main

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"sync"
	"syscall"
	"time"
)

type LockEntry struct {
	Locked bool   `json:"locked"`
	Key    string `json:"key,omitempty"`
	Time   string `json:"time,omitempty"`
}

type LockFile map[string]*LockEntry

// heldLock tracks the meta-lock currently held by this process so it can be
// released on panic, signal, or any other unexpected exit.
var (
	heldMu      sync.Mutex
	heldMgrPath string
	heldMgrKey  string
)

func setHeldLock(path, key string) {
	heldMu.Lock()
	heldMgrPath = path
	heldMgrKey = key
	heldMu.Unlock()
}

func clearHeldLock() {
	heldMu.Lock()
	heldMgrPath = ""
	heldMgrKey = ""
	heldMu.Unlock()
}

// cleanupMgrLock releases the meta-lock if this process still holds it.
// Safe to call multiple times.
func cleanupMgrLock() {
	heldMu.Lock()
	path := heldMgrPath
	key := heldMgrKey
	heldMu.Unlock()

	if path == "" || key == "" {
		return
	}
	lf, err := readLockFile(path)
	if err != nil {
		return
	}
	entry := lf[path]
	if entry != nil && entry.Key == key {
		delete(lf, path)
		_ = writeLockFile(path, lf)
	}
	clearHeldLock()
}

func generateUUID() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	b[6] = (b[6] & 0x0f) | 0x40 // version 4
	b[8] = (b[8] & 0x3f) | 0x80 // variant bits
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]), nil
}

func readLockFile(path string) (LockFile, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return make(LockFile), nil
		}
		return nil, fmt.Errorf("read lock file: %w", err)
	}
	if len(data) == 0 {
		return make(LockFile), nil
	}
	var lf LockFile
	if err := json.Unmarshal(data, &lf); err != nil {
		return make(LockFile), nil
	}
	if lf == nil {
		return make(LockFile), nil
	}
	return lf, nil
}

func writeLockFile(path string, lf LockFile) error {
	data, err := json.MarshalIndent(lf, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal lock file: %w", err)
	}
	return os.WriteFile(path, data, 0644)
}

// acquireMgrLock implements steps 3-8: acquire the meta-lock on the JSON file itself.
// Returns the mgr-key on success.
func acquireMgrLock(mgrPath string) (string, error) {
	// Step 3: generate mgr-key once
	mgrKey, err := generateUUID()
	if err != nil {
		return "", err
	}

	for {
		// Step 4: record current time
		startTime := time.Now()

		// Steps 5-6: spin until the meta-lock is free or owned by us
		for {
			lf, err := readLockFile(mgrPath)
			if err != nil {
				return "", err
			}
			entry := lf[mgrPath]
			if entry == nil || !entry.Locked || entry.Key == mgrKey {
				break
			}
			if time.Since(startTime) > 5*time.Second {
				return "", fmt.Errorf("timeout waiting to acquire manager lock")
			}
			time.Sleep(100 * time.Millisecond)
		}

		// Step 7: write locked=true, key=mgrKey
		lf, err := readLockFile(mgrPath)
		if err != nil {
			return "", err
		}
		if lf[mgrPath] == nil {
			lf[mgrPath] = &LockEntry{}
		}
		lf[mgrPath].Locked = true
		lf[mgrPath].Key = mgrKey
		if err := writeLockFile(mgrPath, lf); err != nil {
			return "", err
		}

		// Step 8: verify we actually won the race
		lf2, err := readLockFile(mgrPath)
		if err != nil {
			return "", err
		}
		if lf2[mgrPath] != nil && lf2[mgrPath].Key == mgrKey {
			setHeldLock(mgrPath, mgrKey)
			return mgrKey, nil
		}
		// Lost the race; go back to step 4 (keep same mgrKey, reset timer)
	}
}

// releaseMgrLock removes the meta-lock entry and writes the file.
func releaseMgrLock(mgrPath string) error {
	lf, err := readLockFile(mgrPath)
	if err != nil {
		return err
	}
	delete(lf, mgrPath)
	err = writeLockFile(mgrPath, lf)
	if err == nil {
		clearHeldLock()
	}
	return err
}

func cmdAcquire(mgrPath, filePath, key string) error {
	// Steps 3-8: acquire meta-lock
	if _, err := acquireMgrLock(mgrPath); err != nil {
		return err
	}

	// Step 9: record start time for file-lock wait timeout
	startTime := time.Now()

	for {
		// Step 10: read current file lock state
		lf, err := readLockFile(mgrPath)
		if err != nil {
			_ = releaseMgrLock(mgrPath)
			return err
		}

		entry := lf[filePath]
		xLocked := entry != nil && entry.Locked
		xKey := ""
		if entry != nil {
			xKey = entry.Key
		}

		// Step 11: if locked by someone else, wait and retry
		if xLocked && xKey != key {
			if time.Since(startTime) > 30*time.Second {
				_ = releaseMgrLock(mgrPath)
				return fmt.Errorf("timeout waiting to acquire lock on %s", filePath)
			}
			// Release meta-lock while sleeping so others can proceed
			if err := releaseMgrLock(mgrPath); err != nil {
				return err
			}
			time.Sleep(1 * time.Second)
			// Re-acquire meta-lock before next read
			if _, err := acquireMgrLock(mgrPath); err != nil {
				return err
			}
			continue
		}

		// Step 12: not locked (or already owned by caller's key) — acquire
		if lf[filePath] == nil {
			lf[filePath] = &LockEntry{}
		}
		lf[filePath].Locked = true
		lf[filePath].Key = key
		lf[filePath].Time = time.Now().UTC().Format(time.RFC3339)

		// Step 13: delete meta-lock entry and write atomically
		delete(lf, mgrPath)
		if err := writeLockFile(mgrPath, lf); err != nil {
			return err
		}
		clearHeldLock()
		return nil
	}
}

func cmdRelease(mgrPath, filePath, key string) error {
	// Steps 3-8: acquire meta-lock
	if _, err := acquireMgrLock(mgrPath); err != nil {
		return err
	}

	lf, err := readLockFile(mgrPath)
	if err != nil {
		_ = releaseMgrLock(mgrPath)
		return err
	}

	entry := lf[filePath]
	xLocked := entry != nil && entry.Locked
	xKey := ""
	if entry != nil {
		xKey = entry.Key
	}

	// Step 14: validate preconditions
	if !xLocked {
		_ = releaseMgrLock(mgrPath)
		return fmt.Errorf("file %s is not locked", filePath)
	}
	if xKey != key {
		_ = releaseMgrLock(mgrPath)
		return fmt.Errorf("key does not match the lock key for %s", filePath)
	}

	// Step 15: delete file lock and meta-lock, write
	delete(lf, filePath)
	delete(lf, mgrPath)
	err = writeLockFile(mgrPath, lf)
	if err == nil {
		clearHeldLock()
	}
	return err
}

func cmdForceUnlock(mgrPath, filePath string) error {
	// Steps 3-8: acquire meta-lock
	if _, err := acquireMgrLock(mgrPath); err != nil {
		return err
	}

	lf, err := readLockFile(mgrPath)
	if err != nil {
		_ = releaseMgrLock(mgrPath)
		return err
	}

	// Remove file lock and meta-lock unconditionally
	delete(lf, filePath)
	delete(lf, mgrPath)
	err = writeLockFile(mgrPath, lf)
	if err == nil {
		clearHeldLock()
	}
	return err
}

func run() error {
	// cleanupMgrLock runs on normal return AND on panic (defer unwinds through panics).
	// os.Exit bypasses defer, so we keep os.Exit only in main() after run() returns.
	defer cleanupMgrLock()

	action := ""
	if len(os.Args) >= 2 {
		action = os.Args[1]
	}

	if action == "force-unlock" && len(os.Args) < 3 {
		return fmt.Errorf("usage: lock-mgr force-unlock <file>")
	}
	if (action == "acquire" || action == "release") && len(os.Args) < 4 {
		return fmt.Errorf("usage: lock-mgr %s <file> <key>", action)
	}
	if action == "" || (action != "acquire" && action != "release" && action != "force-unlock") {
		return fmt.Errorf("usage: lock-mgr <acquire|release> <file> <key>\n       lock-mgr force-unlock <file>")
	}

	fileArg := os.Args[2]
	key := ""
	if action != "force-unlock" {
		key = os.Args[3]
	}

	// Step 1: resolve tool directory and locate (or create) the lock file
	execPath, err := os.Executable()
	if err != nil {
		return fmt.Errorf("cannot determine executable path: %w", err)
	}

	toolDir := filepath.Dir(execPath)
	rawMgrPath := filepath.Join(toolDir, "..", ".lock-mgr.json")

	// Step 2: get absolute paths
	mgrPath, err := filepath.Abs(rawMgrPath)
	if err != nil {
		return fmt.Errorf("cannot resolve lock file path: %w", err)
	}

	filePath, err := filepath.Abs(fileArg)
	if err != nil {
		return fmt.Errorf("cannot resolve file path: %w", err)
	}

	// Ensure lock file exists
	if _, statErr := os.Stat(mgrPath); os.IsNotExist(statErr) {
		if writeErr := os.WriteFile(mgrPath, []byte("{}"), 0644); writeErr != nil {
			return fmt.Errorf("cannot create lock file: %w", writeErr)
		}
	}

	switch action {
	case "acquire":
		return cmdAcquire(mgrPath, filePath, key)
	case "release":
		return cmdRelease(mgrPath, filePath, key)
	case "force-unlock":
		return cmdForceUnlock(mgrPath, filePath)
	}
	return nil
}

func main() {
	// Signal handler: release meta-lock on SIGINT / SIGTERM before exiting.
	// This covers Ctrl+C and process termination while the lock is held.
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		cleanupMgrLock()
		os.Exit(130)
	}()

	if err := run(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
}