iSharkFly-Docs
/
packer-cn
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add WaitGroup to avoid data race

This commit is contained in:
sylviamoss 2020-09-28 15:17:59 +02:00 committed by Wilken Rivera
parent 1c2b469acd
commit b2c7897f58
1 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
builder/amazon/common/ssm_driver.go
View File

@ -7,6 +7,7 @@ import (
"log" "log"
"os/exec" "os/exec"
"strings" "strings"
"sync"
"time" "time"
"github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/aws"
@ -35,7 +36,9 @@ type SSMDriver struct {
session *ssm.StartSessionOutput session *ssm.StartSessionOutput
sessionParams ssm.StartSessionInput sessionParams ssm.StartSessionInput
pluginCmdFunc func(context.Context) error pluginCmdFunc func(context.Context) error
retryConnection chan bool retryConnection chan bool
wg sync.WaitGroup
} }
func NewSSMDriver(config SSMDriverConfig) *SSMDriver { func NewSSMDriver(config SSMDriverConfig) *SSMDriver {
@ -48,6 +51,8 @@ func NewSSMDriver(config SSMDriverConfig) *SSMDriver {
// not wish to manage the session manually calling StopSession on a instance of this driver will terminate the active session // not wish to manage the session manually calling StopSession on a instance of this driver will terminate the active session
// created from calling StartSession. // created from calling StartSession.
func (d *SSMDriver) StartSession(ctx context.Context, input ssm.StartSessionInput) (*ssm.StartSessionOutput, error) { func (d *SSMDriver) StartSession(ctx context.Context, input ssm.StartSessionInput) (*ssm.StartSessionOutput, error) {
d.wg.Add(1)
defer d.wg.Done()
log.Printf("Starting PortForwarding session to instance %q", aws.StringValue(input.Target)) log.Printf("Starting PortForwarding session to instance %q", aws.StringValue(input.Target))
var output *ssm.StartSessionOutput var output *ssm.StartSessionOutput
@ -76,6 +81,7 @@ func (d *SSMDriver) StartSession(ctx context.Context, input ssm.StartSessionInpu
case <-driver.retryConnection: case <-driver.retryConnection:
if retryTimes <= 11 { if retryTimes <= 11 {
retryTimes++ retryTimes++
d.wg.Wait()
_, err := driver.StartSession(ctx, input) _, err := driver.StartSession(ctx, input)
if err != nil { if err != nil {
return return
@ -193,6 +199,9 @@ func isRetryableError(output string) bool {
// StopSession terminates an active Session Manager session // StopSession terminates an active Session Manager session
func (d *SSMDriver) StopSession() error { func (d *SSMDriver) StopSession() error {
d.wg.Add(1)
defer d.wg.Done()
if d.retryConnection != nil { if d.retryConnection != nil {
close(d.retryConnection) close(d.retryConnection)
} }