Machine lease UX improvements (#3815)

* bump fly-go for lease context fix

* machine lease improvements

- print status update if we're still waiting on lease after 2s
- release leases even  if command has been aborted

* improve deploy log messages around lease acquisition

* skip acquiring nonce if we already have it

* fix panic in postgres import

Author: btoews

go.mod

@@ -72,7 +72,7 @@ require (
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.9.0
-	github.com/superfly/fly-go v0.1.21
+	github.com/superfly/fly-go v0.1.22-0.20240802153530-1487eea07345
 	github.com/superfly/graphql v0.2.4
 	github.com/superfly/lfsc-go v0.1.1
 	github.com/superfly/macaroon v0.2.14-0.20240702184853-b8ac52a1fc77
@@ -272,5 +272,5 @@ require (
 replace (
 	github.com/loadsmart/calver-go => github.com/ndarilek/calver-go v0.0.0-20230710153822-893bbd83a936
 	github.com/nats-io/nats.go => github.com/btoews/nats.go v0.0.0-20240401180931-476bea7f4158
-	github.com/superfly/fly-go => github.com/superfly/fly-go v0.1.21
+	github.com/superfly/fly-go => github.com/superfly/fly-go v0.1.22-0.20240802153530-1487eea07345
 )

go.sum

@@ -629,8 +629,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8=
 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
-github.com/superfly/fly-go v0.1.21 h1:+UEmy+4jZ0f/lUWxF1u/y54qbodOCqcrvuL92ANMePo=
-github.com/superfly/fly-go v0.1.21/go.mod h1:JQke/BwoZqrWurqYkypSlcSo7bIUgCI3eVnqMC6AUj0=
+github.com/superfly/fly-go v0.1.22-0.20240802153530-1487eea07345 h1:4/khSPMwvoXgTyj8cKTu54dSFuIPkK0sTrMysTYa7+4=
+github.com/superfly/fly-go v0.1.22-0.20240802153530-1487eea07345/go.mod h1:JQke/BwoZqrWurqYkypSlcSo7bIUgCI3eVnqMC6AUj0=
 github.com/superfly/graphql v0.2.4 h1:Av8hSk4x8WvKJ6MTnEwrLknSVSGPc7DWpgT3z/kt3PU=
 github.com/superfly/graphql v0.2.4/go.mod h1:CVfDl31srm8HnJ9udwLu6hFNUW/P6GUM2dKcG1YQ8jc=
 github.com/superfly/lfsc-go v0.1.1 h1:dGjLgt81D09cG+aR9lJZIdmonjZSR5zYCi7s54+ZU2Q=

internal/command/deploy/plan.go

@@ -321,12 +321,14 @@ func (md *machineDeployment) acquireLeases(ctx context.Context, machineTuples []
 			} else {
 				return nil
 			}
-			sl.LogStatus(statuslogger.StatusRunning, fmt.Sprintf("Acquiring lease for %s", machine.ID))
 
-			if machine.LeaseNonce == "" {
-				sl.LogStatus(statuslogger.StatusRunning, fmt.Sprintf("Waiting for job %s", machine.ID))
+			if machine.LeaseNonce != "" {
+				sl.LogStatus(statuslogger.StatusRunning, fmt.Sprintf("Already have lease for %s", machine.ID))
+				return nil
 			}
 
+			sl.LogStatus(statuslogger.StatusRunning, fmt.Sprintf("Acquiring lease for %s", machine.ID))
+
 			lease, err := md.acquireMachineLease(ctx, machine.ID)
 			if err != nil {
 				sl.LogStatus(statuslogger.StatusFailure, fmt.Sprintf("Failed to acquire lease for %s: %v", machine.ID, err))
@@ -336,7 +338,7 @@ func (md *machineDeployment) acquireLeases(ctx context.Context, machineTuples []
 			machine.LeaseNonce = lease.Data.Nonce
 			lm := mach.NewLeasableMachine(md.flapsClient, md.io, machine, false)
 			lm.StartBackgroundLeaseRefresh(ctx, md.leaseTimeout, md.leaseDelayBetween)
-			sl.LogStatus(statuslogger.StatusRunning, fmt.Sprintf("Waiting for job %s", machine.ID))
+			sl.LogStatus(statuslogger.StatusRunning, fmt.Sprintf("Acquired lease for %s", machine.ID))
 			return nil
 		})
 	}

internal/command/postgres/import.go

@@ -115,6 +115,9 @@ func runImport(ctx context.Context) error {
 		return fmt.Errorf("no machines are available on this app %s", appName)
 	}
 	leader, _ := machinesNodeRoles(ctx, machines)
+	if leader == nil {
+		return fmt.Errorf("no active leader found")
+	}
 	machineID := leader.ID
 
 	// Resolve region

internal/machine/lease.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/sourcegraph/conc/pool"
 	fly "github.com/superfly/fly-go"
@@ -60,16 +61,41 @@ func releaseLease(ctx context.Context, machine *fly.Machine) {
 	io := iostreams.FromContext(ctx)
 	flapsClient := flapsutil.ClientFromContext(ctx)
 
+	// remove the cancel from ctx so we can still releases leases if the command was aborted
+	if ctx.Err() != nil {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
+		defer cancel()
+
+		fmt.Fprintf(io.Out, "Releasing lease for machine %s...\n", machine.ID)
+	}
+
 	if err := flapsClient.ReleaseLease(ctx, machine.ID, machine.LeaseNonce); err != nil {
 		if !strings.Contains(err.Error(), "lease not found") {
-			fmt.Fprintf(io.Out, "failed to release lease for machine %s: %s", machine.ID, err.Error())
+			fmt.Fprintf(io.Out, "failed to release lease for machine %s: %s\n", machine.ID, err.Error())
 		}
 	}
 }
 
 // AcquireLease works to acquire/attach a lease for the specified machine.
 // WARNING: Make sure you defer the lease release process.
 func AcquireLease(ctx context.Context, machine *fly.Machine) (*fly.Machine, releaseLeaseFunc, error) {
+	// if we haven't gotten the lease after 2s, we print a message so users
+	// aren't left wondering.
+	abortStatusUpdate := make(chan struct{})
+	defer close(abortStatusUpdate)
+	go func() {
+		timer := time.NewTimer(2 * time.Second)
+		defer timer.Stop()
+
+		select {
+		case <-timer.C:
+			io := iostreams.FromContext(ctx)
+			fmt.Fprintf(io.Out, "Waiting on lease for machine %s...\n", machine.ID)
+		case <-abortStatusUpdate:
+		}
+	}()
+
 	flapsClient := flapsutil.ClientFromContext(ctx)
 
 	lease, err := flapsClient.AcquireLease(ctx, machine.ID, fly.IntPointer(120))