support custom lock expiration

README.md

@@ -69,4 +69,4 @@ $ go run main.go run testkey 'echo "hello world"' -t loomctl-locks
 {"level":"info","ts":1651558192.944402,"caller":"cmd/run.go:33","msg":"command succeeded","key":"testkey","command":"echo \"hello world\"","output":"\"hello world\"\n"}
 ```
 
-Once the command finishes running `lock-exec` will unlock the key. It also listens for os interrupts and unlocks the key before exiting. In the rare case where `lock-exec` exits and fails to unlock the key will remain locked for the next 24 hours. The key can manually be unlocked earlier using `lock-exec unlock <key>`.
+Once the command finishes running `lock-exec` will unlock the key. It also listens for os interrupts and unlocks the key before exiting. In the rare case where `lock-exec` exits and fails to unlock the key will remain locked for the next 24 hours (you can customize this with `--expire <duration>`). The key can manually be unlocked earlier using `lock-exec unlock <key>`.

cmd/cli.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
@@ -22,6 +23,7 @@ type cli struct {
 	log *zap.SugaredLogger
 
 	table   string
+	expire  time.Duration
 	version string
 }
 

cmd/root.go

@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"time"
+
 	"github.com/spf13/cobra"
 )
 
@@ -15,6 +17,7 @@ func (c *cli) newRootCmd() *cobra.Command {
 	}
 
 	cmd.PersistentFlags().StringVarP(&c.table, "table", "t", "lock-exec", "table name in dynamodb to use for locking")
+	cmd.PersistentFlags().DurationVarP(&c.expire, "expire", "e", time.Hour*24, "lock duration in the event that the post-run unlock fails") //nolint:gomnd
 
 	cmd.AddCommand(
 		c.newRunCmd(),

cmd/run.go

@@ -18,6 +18,7 @@ func (c *cli) newRunCmd() *cobra.Command {
 		Run: func(cmd *cobra.Command, args []string) {
 			log := c.log.With("key", args[0], "command", args[1])
 			locker := c.newLocker()
+			locker.SetExpire(c.expire)
 
 			log.Info("running command")
 			err := locker.Run(c.cmd.Context(), args[0], args[1])

lock/client.go

@@ -3,6 +3,7 @@ package lock
 import (
 	"context"
 	"errors"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
 )
@@ -23,9 +24,19 @@ type storageI interface {
 type Client struct {
 	storage storageI
 	table   string
+	expire  time.Duration
 }
 
 // New creates a new lock client.
 func New(ddb storageI, table string) *Client {
-	return &Client{storage: ddb, table: table}
+	return &Client{
+		storage: ddb,
+		table:   table,
+		expire:  time.Hour * 24, //nolint:gomnd
+	}
+}
+
+// SetExpire sets the default expire duration for locks. Defaults to 24 hours if not set.
+func (c *Client) SetExpire(d time.Duration) {
+	c.expire = d
 }

lock/run.go

@@ -6,7 +6,6 @@ import (
 	"os"
 	"os/exec"
 	"strings"
-	"time"
 )
 
 // Run acquires a lock under the specified key, executes the command, and then unlocks the key.
@@ -16,7 +15,7 @@ func (c *Client) Run(ctx context.Context, key, command string) error {
 	// use context.Background here so that unlock runs even if the context is cancelled
 	defer c.Unlock(context.Background(), key) //nolint:errcheck,contextcheck
 
-	err := c.Lock(ctx, key, time.Hour*24) //nolint:gomnd
+	err := c.Lock(ctx, key, c.expire)
 	if err != nil {
 		return fmt.Errorf("lock failed: %w", err)
 	}