From 19e8f150a3f9b8ecd16fc52bef5ec9910e4f9df0 Mon Sep 17 00:00:00 2001 From: Vladislav Rassokhin Date: Fri, 12 Jul 2019 12:54:43 +0300 Subject: [PATCH] Use context for timeouts, interruption in ssh and winrm communicators Also don't waste 5 seconds waiting before first winrm connection attempt Minor code cleanup as well --- helper/communicator/step_connect_ssh.go | 29 +++++++------- helper/communicator/step_connect_winrm.go | 47 ++++++++++++----------- 2 files changed, 37 insertions(+), 39 deletions(-) diff --git a/helper/communicator/step_connect_ssh.go b/helper/communicator/step_connect_ssh.go index d2d1789fd..28347fe95 100644 --- a/helper/communicator/step_connect_ssh.go +++ b/helper/communicator/step_connect_ssh.go @@ -36,17 +36,17 @@ func (s *StepConnectSSH) Run(ctx context.Context, state multistep.StateBag) mult var comm packer.Communicator var err error - cancel := make(chan struct{}) + subCtx, cancel := context.WithCancel(ctx) waitDone := make(chan bool, 1) go func() { ui.Say("Waiting for SSH to become available...") - comm, err = s.waitForSSH(state, cancel) + comm, err = s.waitForSSH(state, subCtx) + cancel() // just to make 'possible context leak' analysis happy waitDone <- true }() log.Printf("[INFO] Waiting for SSH, up to timeout: %s", s.Config.SSHTimeout) timeout := time.After(s.Config.SSHTimeout) -WaitLoop: for { // Wait for either SSH to become available, a timeout to occur, // or an interrupt to come through. @@ -60,31 +60,28 @@ WaitLoop: ui.Say("Connected to SSH!") state.Put("communicator", comm) - break WaitLoop + return multistep.ActionContinue case <-timeout: err := fmt.Errorf("Timeout waiting for SSH.") state.Put("error", err) ui.Error(err.Error()) - close(cancel) + cancel() + return multistep.ActionHalt + case <-ctx.Done(): + // The step sequence was cancelled, so cancel waiting for SSH + // and just start the halting process. + cancel() + log.Println("[WARN] Interrupt detected, quitting waiting for SSH.") return multistep.ActionHalt case <-time.After(1 * time.Second): - if _, ok := state.GetOk(multistep.StateCancelled); ok { - // The step sequence was cancelled, so cancel waiting for SSH - // and just start the halting process. - close(cancel) - log.Println("[WARN] Interrupt detected, quitting waiting for SSH.") - return multistep.ActionHalt - } } } - - return multistep.ActionContinue } func (s *StepConnectSSH) Cleanup(multistep.StateBag) { } -func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, cancel <-chan struct{}) (packer.Communicator, error) { +func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, ctx context.Context) (packer.Communicator, error) { // Determine if we're using a bastion host, and if so, retrieve // that configuration. This configuration doesn't change so we // do this one before entering the retry loop. @@ -123,7 +120,7 @@ func (s *StepConnectSSH) waitForSSH(state multistep.StateBag, cancel <-chan stru // Don't check for cancel or wait on first iteration if !first { select { - case <-cancel: + case <-ctx.Done(): log.Println("[DEBUG] SSH wait cancelled. Exiting loop.") return nil, errors.New("SSH wait cancelled") case <-time.After(5 * time.Second): diff --git a/helper/communicator/step_connect_winrm.go b/helper/communicator/step_connect_winrm.go index c9a748b0f..f1e30820a 100644 --- a/helper/communicator/step_connect_winrm.go +++ b/helper/communicator/step_connect_winrm.go @@ -33,23 +33,23 @@ type StepConnectWinRM struct { WinRMPort func(multistep.StateBag) (int, error) } -func (s *StepConnectWinRM) Run(_ context.Context, state multistep.StateBag) multistep.StepAction { +func (s *StepConnectWinRM) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction { ui := state.Get("ui").(packer.Ui) var comm packer.Communicator var err error - cancel := make(chan struct{}) + subCtx, cancel := context.WithCancel(ctx) waitDone := make(chan bool, 1) go func() { ui.Say("Waiting for WinRM to become available...") - comm, err = s.waitForWinRM(state, cancel) + comm, err = s.waitForWinRM(state, subCtx) + cancel() // just to make 'possible context leak' analysis happy waitDone <- true }() log.Printf("Waiting for WinRM, up to timeout: %s", s.Config.WinRMTimeout) timeout := time.After(s.Config.WinRMTimeout) -WaitLoop: for { // Wait for either WinRM to become available, a timeout to occur, // or an interrupt to come through. @@ -62,39 +62,40 @@ WaitLoop: ui.Say("Connected to WinRM!") state.Put("communicator", comm) - break WaitLoop + return multistep.ActionContinue case <-timeout: err := fmt.Errorf("Timeout waiting for WinRM.") state.Put("error", err) ui.Error(err.Error()) - close(cancel) + cancel() + return multistep.ActionHalt + case <-ctx.Done(): + // The step sequence was cancelled, so cancel waiting for WinRM + // and just start the halting process. + cancel() + log.Println("Interrupt detected, quitting waiting for WinRM.") return multistep.ActionHalt case <-time.After(1 * time.Second): - if _, ok := state.GetOk(multistep.StateCancelled); ok { - // The step sequence was cancelled, so cancel waiting for WinRM - // and just start the halting process. - close(cancel) - log.Println("Interrupt detected, quitting waiting for WinRM.") - return multistep.ActionHalt - } } } - - return multistep.ActionContinue } func (s *StepConnectWinRM) Cleanup(multistep.StateBag) { } -func (s *StepConnectWinRM) waitForWinRM(state multistep.StateBag, cancel <-chan struct{}) (packer.Communicator, error) { - ctx := context.TODO() +func (s *StepConnectWinRM) waitForWinRM(state multistep.StateBag, ctx context.Context) (packer.Communicator, error) { var comm packer.Communicator + first := true for { - select { - case <-cancel: - log.Println("[INFO] WinRM wait cancelled. Exiting loop.") - return nil, errors.New("WinRM wait cancelled") - case <-time.After(5 * time.Second): + // Don't check for cancel or wait on first iteration + if !first { + select { + case <-ctx.Done(): + log.Println("[INFO] WinRM wait cancelled. Exiting loop.") + return nil, errors.New("WinRM wait cancelled") + case <-time.After(5 * time.Second): + } + first = false } host, err := s.Host(state) @@ -157,7 +158,7 @@ func (s *StepConnectWinRM) waitForWinRM(state multistep.StateBag, cancel <-chan cmd.Stdout = &buf cmd.Stdout = io.MultiWriter(cmd.Stdout, &buf2) select { - case <-cancel: + case <-ctx.Done(): log.Println("WinRM wait canceled, exiting loop") return comm, fmt.Errorf("WinRM wait canceled") case <-time.After(retryableSleep):