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

     1  package gclient
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"strings"
     7  
     8  	"code.cloudfoundry.org/garden"
     9  	"github.com/pf-qiu/concourse/v6/atc/worker/gclient/connection"
    10  )
    11  
    12  type RetryableConnection struct {
    13  	connection.Connection
    14  }
    15  
    16  func NewRetryableConnection(connection connection.Connection) connection.Connection {
    17  	return &RetryableConnection{
    18  		Connection: connection,
    19  	}
    20  }
    21  
    22  func (conn *RetryableConnection) Run(ctx context.Context, handle string, processSpec garden.ProcessSpec, processIO garden.ProcessIO) (garden.Process, error) {
    23  	innerProcess, err := conn.Connection.Run(ctx, handle, processSpec, processIO)
    24  	if err != nil {
    25  		return nil, err
    26  	}
    27  
    28  	return &retryableProcess{
    29  		Process: innerProcess,
    30  
    31  		rehydrate: func() (garden.Process, error) {
    32  			return conn.Connection.Attach(ctx, handle, innerProcess.ID(), processIO)
    33  		},
    34  	}, nil
    35  }
    36  
    37  func (conn *RetryableConnection) Attach(ctx context.Context, handle string, processID string, processIO garden.ProcessIO) (garden.Process, error) {
    38  	innerProcess, err := conn.Connection.Attach(ctx, handle, processID, processIO)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  
    43  	return &retryableProcess{
    44  		Process: innerProcess,
    45  
    46  		rehydrate: func() (garden.Process, error) {
    47  			return conn.Connection.Attach(ctx, handle, processID, processIO)
    48  		},
    49  	}, nil
    50  }
    51  
    52  type retryableProcess struct {
    53  	garden.Process
    54  
    55  	rehydrate func() (garden.Process, error)
    56  }
    57  
    58  func (process *retryableProcess) Wait() (int, error) {
    59  	for {
    60  		status, err := process.Process.Wait()
    61  		if err == nil {
    62  			return status, nil
    63  		}
    64  
    65  		if !strings.HasSuffix(err.Error(), io.EOF.Error()) {
    66  			return 0, err
    67  		}
    68  
    69  		process.Process, err = process.rehydrate()
    70  		if err != nil {
    71  			return 0, err
    72  		}
    73  	}
    74  }
    75  
    76  func (process *retryableProcess) Signal(sig garden.Signal) error {
    77  	for {
    78  		err := process.Process.Signal(sig)
    79  		if err == nil {
    80  			return nil
    81  		}
    82  
    83  		if strings.Contains(err.Error(), "use of closed network connection") {
    84  			return err
    85  		}
    86  
    87  		process.Process, err = process.rehydrate()
    88  		if err != nil {
    89  			return err
    90  		}
    91  	}
    92  }
    93  
    94  func (process *retryableProcess) SetTTY(tty garden.TTYSpec) error {
    95  	for {
    96  		err := process.Process.SetTTY(tty)
    97  		if err == nil {
    98  			return nil
    99  		}
   100  
   101  		if strings.Contains(err.Error(), "use of closed network connection") {
   102  			return err
   103  		}
   104  
   105  		process.Process, err = process.rehydrate()
   106  		if err != nil {
   107  			return err
   108  		}
   109  	}
   110  }