Merge pull request #652 from erikwilson/update-tunnel-wait

Fix tunnel endpoint upgrade
This commit is contained in:
Erik Wilson 2019-07-17 18:39:45 -07:00 committed by GitHub
commit 75b8d71e77
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 18 additions and 3 deletions

View File

@ -47,7 +47,7 @@ func run(ctx context.Context, cfg cmds.Agent) error {
return err
}
if err := tunnel.Setup(nodeConfig); err != nil {
if err := tunnel.Setup(ctx, nodeConfig); err != nil {
return err
}

View File

@ -52,7 +52,7 @@ func getAddresses(endpoint *v1.Endpoints) []string {
return serverAddresses
}
func Setup(config *config.Node) error {
func Setup(ctx context.Context, config *config.Node) error {
restConfig, err := clientcmd.BuildConfigFromFlags("", config.AgentConfig.KubeConfigNode)
if err != nil {
return err
@ -83,7 +83,6 @@ func Setup(config *config.Node) error {
disconnect[address] = connect(wg, address, config, transportConfig)
}
}
wg.Wait()
go func() {
connect:
@ -134,6 +133,19 @@ func Setup(config *config.Node) error {
}
}()
wait := make(chan int, 1)
go func() {
wg.Wait()
wait <- 0
}()
select {
case <-ctx.Done():
logrus.Error("tunnel context canceled while waiting for connection")
return ctx.Err()
case <-wait:
}
return nil
}
@ -178,6 +190,9 @@ func connect(waitGroup *sync.WaitGroup, address string, config *config.Node, tra
})
if ctx.Err() != nil {
if waitGroup != nil {
once.Do(waitGroup.Done)
}
logrus.Infof("Stopped tunnel to %s", wsURL)
return
}