golang.org/x/build@v0.0.0-20240506185731-218518f32b70/cmd/runqemubuildlet/heartbeat.go (about)

     1  // Copyright 2021 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"io"
    11  	"net/http"
    12  	"time"
    13  
    14  	"golang.org/x/build/internal"
    15  )
    16  
    17  // buildletHealthTimeout is the maximum time to wait for a
    18  // checkBuildletHealth request to complete.
    19  const buildletHealthTimeout = 10 * time.Second
    20  
    21  // checkBuildletHealth performs a GET request against URL, and returns
    22  // an error if an http.StatusOK isn't returned before
    23  // buildletHealthTimeout has elapsed.
    24  func checkBuildletHealth(ctx context.Context, url string) error {
    25  	ctx, cancel := context.WithTimeout(ctx, buildletHealthTimeout)
    26  	defer cancel()
    27  	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
    28  	if err != nil {
    29  		return err
    30  	}
    31  	resp, err := http.DefaultClient.Do(req)
    32  	if err != nil {
    33  		return err
    34  	}
    35  	defer resp.Body.Close()
    36  	if _, err := io.Copy(io.Discard, resp.Body); err != nil {
    37  		return err
    38  	}
    39  	if resp.StatusCode != http.StatusOK {
    40  		return fmt.Errorf("resp.StatusCode = %d, wanted %d", resp.StatusCode, http.StatusOK)
    41  	}
    42  	return nil
    43  }
    44  
    45  // heartbeatContext calls f every period. If f consistently returns an
    46  // error for longer than the provided timeout duration, the context
    47  // returned by heartbeatContext will be cancelled, and
    48  // heartbeatContext will stop sending requests.
    49  //
    50  // A single call to f that does not return an error will reset the
    51  // timeout window, unless heartbeatContext has already timed out.
    52  func heartbeatContext(ctx context.Context, period time.Duration, timeout time.Duration, f func(context.Context) error) (context.Context, func()) {
    53  	ctx, cancel := context.WithCancel(ctx)
    54  
    55  	lastSuccess := time.Now()
    56  	go internal.PeriodicallyDo(ctx, period, func(ctx context.Context, t time.Time) {
    57  		err := f(ctx)
    58  		if err != nil && t.Sub(lastSuccess) > timeout {
    59  			cancel()
    60  		}
    61  		if err == nil {
    62  			lastSuccess = t
    63  		}
    64  	})
    65  
    66  	return ctx, cancel
    67  }