@@ -161,22 +161,7 @@ func main() {
161
161
}
162
162
163
163
if launchCtx .warmup {
164
- ctx , cancel := context .WithCancel (context .Background ())
165
-
166
- out:
167
- for {
168
- log .Debug ("launcher: wait for tasks to finish before running warmup" )
169
- err := waitForTasksToFinish (ctx , cancel )
170
- if err != nil {
171
- log .WithError (err ).Warn ("launcher: failed to observe tasks completion" )
172
- }
173
- select {
174
- case <- ctx .Done ():
175
- break out
176
- case <- time .After (1 * time .Second ):
177
- }
178
- }
179
-
164
+ starWarmup ()
180
165
launch (launchCtx )
181
166
return
182
167
}
@@ -187,6 +172,31 @@ func main() {
187
172
serve (launchCtx )
188
173
}
189
174
175
+ func starWarmup () {
176
+ ctx , cancel := context .WithCancel (context .Background ())
177
+ var conn * grpc.ClientConn
178
+ var err error
179
+
180
+ for {
181
+ conn , err = dial (ctx )
182
+ if err == nil {
183
+ log .Debug ("launcher: wait for tasks to finish before running warmup" )
184
+ finished , err := waitForTasksToFinish (ctx , conn )
185
+ if err != nil {
186
+ log .WithError (err ).Warn ("launcher: failed to observe tasks completion" )
187
+ }
188
+ if finished {
189
+ cancel ()
190
+ }
191
+ }
192
+ select {
193
+ case <- ctx .Done ():
194
+ return
195
+ case <- time .After (1 * time .Second ):
196
+ }
197
+ }
198
+ }
199
+
190
200
func serve (launchCtx * LaunchContext ) {
191
201
debugAgentPrefix := "-agentlib:jdwp=transport=dt_socket,server=y,suspend=n,address=*:"
192
202
http .HandleFunc ("/debug" , func (w http.ResponseWriter , r * http.Request ) {
@@ -870,20 +880,16 @@ func resolveProjectContextDir(launchCtx *LaunchContext) string {
870
880
return launchCtx .projectDir
871
881
}
872
882
873
- func waitForTasksToFinish (ctx context.Context , cancel context.CancelFunc ) error {
874
- conn , err := dial (ctx )
875
- if err != nil {
876
- return err
877
- }
883
+ func waitForTasksToFinish (ctx context.Context , conn * grpc.ClientConn ) (bool , error ) {
878
884
client := supervisor .NewStatusServiceClient (conn )
879
885
tasksResponse , err := client .TasksStatus (ctx , & supervisor.TasksStatusRequest {Observe : true })
880
886
if err != nil {
881
- return xerrors .Errorf ("failed get tasks status client: %w" , err )
887
+ return false , xerrors .Errorf ("failed get tasks status client: %w" , err )
882
888
}
883
889
for {
884
890
resp , err := tasksResponse .Recv ()
885
891
if err != nil {
886
- return err
892
+ return false , err
887
893
}
888
894
889
895
var runningTasksCounter int
@@ -893,12 +899,10 @@ func waitForTasksToFinish(ctx context.Context, cancel context.CancelFunc) error
893
899
}
894
900
}
895
901
if runningTasksCounter == 0 {
896
- cancel ()
902
+ return true , nil
897
903
}
898
-
899
- return nil
900
904
}
901
-
905
+ return false , nil
902
906
}
903
907
904
908
func dial (ctx context.Context ) (* grpc.ClientConn , error ) {
0 commit comments