@@ -27,7 +27,6 @@ import (
2727 "io"
2828 "math"
2929 "strings"
30- "sync"
3130 "time"
3231
3332 "github.com/gofrs/uuid"
@@ -288,8 +287,6 @@ func evaluateHealthCheckResult(output string) error {
288287 if err != nil {
289288 return err
290289 }
291- log .Warnf ("Health check output: %v\n " , jsonData )
292-
293290 if v , ok := jsonData [headersInstalledPercColumn ]; ok {
294291 switch t := v .(type ) {
295292 case float64 :
@@ -304,90 +301,81 @@ func evaluateHealthCheckResult(output string) error {
304301 return nil
305302}
306303
304+ type healthCheckData struct {
305+ line string
306+ err error
307+ }
308+
309+ // Runs the health check script on the specified vizier. The script's output is evaluated with
310+ // the evaluateHealthCheckResult function to determine if the cluster is healthy. Only a single
311+ // line of output will be parsed from the script.
307312func runHealthCheckScript (v * Connector , execScript * script.ExecutableScript ) (chan string , error ) {
308- output := make (chan string , 1 )
309313 ctx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
310314 defer cancel ()
311315
312316 var encOpts , decOpts * vizierpb.ExecuteScriptRequest_EncryptionOptions
313317
314318 resp , err := RunScript (ctx , []* Connector {v }, execScript , encOpts )
315319 if err != nil {
316- return output , err
320+ return nil , err
317321 }
318322
319323 reader , writer := io .Pipe ()
324+ defer writer .Close ()
320325 defer reader .Close ()
321326 factoryFunc := func (md * vizierpb.ExecuteScriptResponse_MetaData ) components.OutputStreamWriter {
322327 return components .CreateStreamWriter ("json" , writer )
323328 }
324329 tw := NewStreamOutputAdapterWithFactory (ctx , resp , "json" , decOpts , factoryFunc )
325330
326331 bufReader := bufio .NewReader (reader )
327- errCh := make (chan error , 2 )
328- var wg sync. WaitGroup
329- wg . Add ( 1 )
332+ errCh := make (chan error , 1 )
333+ streamCh := make ( chan healthCheckData , 1 )
334+ outputCh := make ( chan string , 1 )
330335 go func () {
331- defer wg .Done ()
332- defer writer .Close ()
333- err = tw .WaitForCompletion ()
334-
335- if err != nil {
336- log .Warnf ("Error on tw.WaitForCompletion: %v" , err )
337- errCh <- err
338- return
336+ defer close (streamCh )
337+ for {
338+ line , err := bufReader .ReadString ('\n' )
339+ streamCh <- healthCheckData {line , err }
340+ if err != nil {
341+ return
342+ }
339343 }
340344 }()
341-
342- wg .Add (1 )
345+ // Consumes the first line of output from the stream or the error from the context.
346+ // The px/agent_status_diagnostics script only outputs one line, but even in the case
347+ // that the fallback (px/agent_status) is used, a single line informs whether the output
348+ // can be processed properly.
343349 go func () {
344- defer wg .Done ()
345- defer close (output )
346- var prevLine string
350+ defer close (errCh )
351+ defer close (outputCh )
347352 for {
348353 select {
349354 case <- ctx .Done ():
350355 errCh <- ctx .Err ()
351356 return
352- default :
353- if ctx .Err () != nil {
354- errCh <- ctx .Err ()
355- return
356- }
357- line , err := bufReader .ReadString ('\n' )
358- if err != nil {
359- if err == io .EOF {
360- log .Warn ("EOF reached while reading health check script output" )
361- output <- prevLine
362- errCh <- err
363- return
364- }
365- errCh <- err
366- return
367- }
368- // Capture the last line of output. This ensures
369- // that the EOF case returns the actual output instead of
370- // an EOF string.
371- prevLine = line
372- err = evaluateHealthCheckResult (line )
373- if err != nil {
374- log .Warn ("evaluateHealthCheckResult err" )
375- output <- prevLine
376- errCh <- err
377- return
357+ case data := <- streamCh :
358+ line := data .line
359+ err := data .err
360+ if err == nil {
361+ err = evaluateHealthCheckResult (line )
378362 }
363+ outputCh <- line
364+ errCh <- err
365+ return
366+
379367 }
380368 }
381369 }()
382370
383- go func () {
384- wg .Wait ()
385- close (errCh )
386- }()
371+ err = tw .WaitForCompletion ()
387372
373+ if err != nil {
374+ return outputCh , err
375+ }
388376 err = <- errCh
389377
390- return output , err
378+ return outputCh , err
391379}
392380
393381// RunSimpleHealthCheckScript runs a diagnostic pxl script to verify query serving works.
0 commit comments