github.com/pf-qiu/concourse/v6@v6.7.3-0.20201207032516-1f455d73275f/atc/worker/transport/hijackable_client.go (about)

     1  package transport
     2  
     3  import (
     4  	"net/http"
     5  
     6  	"github.com/concourse/retryhttp"
     7  )
     8  
     9  type hijackableClient struct {
    10  	db                    TransportDB
    11  	workerName            string
    12  	innerHijackableClient retryhttp.HijackableClient
    13  	cachedHost            *string
    14  }
    15  
    16  func NewHijackableClient(workerName string, db TransportDB, innerHijackableClient retryhttp.HijackableClient) retryhttp.HijackableClient {
    17  	return &hijackableClient{
    18  		innerHijackableClient: innerHijackableClient,
    19  		workerName:            workerName,
    20  		db:                    db,
    21  		cachedHost:            nil,
    22  	}
    23  }
    24  
    25  func (c *hijackableClient) Do(request *http.Request) (*http.Response, retryhttp.HijackCloser, error) {
    26  	if c.cachedHost == nil {
    27  		savedWorker, found, err := c.db.GetWorker(c.workerName)
    28  		if err != nil {
    29  			return nil, nil, err
    30  		}
    31  
    32  		if !found {
    33  			return nil, nil, WorkerMissingError{WorkerName: c.workerName}
    34  		}
    35  
    36  		if savedWorker.GardenAddr() == nil {
    37  			return nil, nil, WorkerUnreachableError{
    38  				WorkerName:  c.workerName,
    39  				WorkerState: string(savedWorker.State()),
    40  			}
    41  		}
    42  
    43  		c.cachedHost = savedWorker.GardenAddr()
    44  	}
    45  
    46  	updatedURL := *request.URL
    47  	updatedURL.Host = *c.cachedHost
    48  
    49  	updatedRequest := *request
    50  	updatedRequest.URL = &updatedURL
    51  
    52  	response, hijackCloser, err := c.innerHijackableClient.Do(&updatedRequest)
    53  	if err != nil {
    54  		c.cachedHost = nil
    55  	}
    56  	return response, hijackCloser, err
    57  }