refactor async lambda
@@ -433,8 +433,7 @@ func (v *Vibranium) RunAndWait(stream pb.CoreRPC_RunAndWaitServer) error {
return err
}
- ctx := stream.Context()
- cancel := context.CancelFunc(func() {})
+ ctx, cancel := context.WithCancel(stream.Context())
if RunAndWaitOptions.Async {
timeout := v.config.GlobalTimeout
if RunAndWaitOptions.AsyncTimeout != 0 {
@@ -464,29 +463,32 @@ func (v *Vibranium) RunAndWait(stream pb.CoreRPC_RunAndWaitServer) error {
}
}()
+ runAndWait := func(f func(<-chan *types.AttachContainerMessage)) error {
+ defer v.taskDone("RunAndWait", true)
+ defer cancel()
ch, err := v.cluster.RunAndWait(ctx, deployOpts, inCh)
if err != nil {
log.Errorf("[RunAndWait] Start run and wait failed %s", err)
- v.taskDone("RunAndWait", true)
return err
}
+ f(ch)
+ return nil
+ }
- sync := make(chan struct{})
- go func() {
- defer v.taskDone("RunAndWait", true)
- defer cancel()
- defer close(sync)
+ if !RunAndWaitOptions.Async {
+ return runAndWait(func(ch <-chan *types.AttachContainerMessage) {
for m := range ch {
- if RunAndWaitOptions.Async {
- log.Infof("[Async RunAndWait] %v", string(m.Data))
- } else if err = stream.Send(toRPCAttachContainerMessage(m)); err != nil {
+ if err = stream.Send(toRPCAttachContainerMessage(m)); err != nil {
v.logUnsentMessages("RunAndWait", m)
}
}
- }()
- if !RunAndWaitOptions.Async {
- <-sync
+ })
}
+ go runAndWait(func(ch <-chan *types.AttachContainerMessage) {
+ for m := range ch {
+ log.Infof("[Async RunAndWait] %v", string(m.Data))
+ }
+ })
return nil
})
}
