github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/helper/pluginutils/grpcutils/utils.go (about)

     1  package grpcutils
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/hashicorp/nomad/plugins/base/structs"
     8  	"google.golang.org/grpc/codes"
     9  	"google.golang.org/grpc/status"
    10  )
    11  
    12  // HandleReqCtxGrpcErr is used to handle a non io.EOF error in a GRPC request
    13  // where a user supplied context is used. It handles detecting if the plugin has
    14  // shutdown via the passeed pluginCtx. The parameters are:
    15  // - err: the error returned from the streaming RPC
    16  // - reqCtx: the user context passed to the request
    17  // - pluginCtx: the plugins done ctx used to detect the plugin dying
    18  //
    19  // The return values are:
    20  // - ErrPluginShutdown if the error is because the plugin shutdown
    21  // - context.Canceled if the reqCtx is canceled
    22  // - The original error
    23  func HandleReqCtxGrpcErr(err error, reqCtx, pluginCtx context.Context) error {
    24  	if err == nil {
    25  		return nil
    26  	}
    27  
    28  	// Determine if the error is because the plugin shutdown
    29  	if errStatus, ok := status.FromError(err); ok &&
    30  		(errStatus.Code() == codes.Unavailable || errStatus.Code() == codes.Canceled) {
    31  		// Potentially wait a little before returning an error so we can detect
    32  		// the exit
    33  		select {
    34  		case <-pluginCtx.Done():
    35  			err = structs.ErrPluginShutdown
    36  		case <-reqCtx.Done():
    37  			err = reqCtx.Err()
    38  
    39  			// There is no guarantee that the select will choose the
    40  			// doneCtx first so we have to double check
    41  			select {
    42  			case <-pluginCtx.Done():
    43  				err = structs.ErrPluginShutdown
    44  			default:
    45  			}
    46  		case <-time.After(3 * time.Second):
    47  			// Its okay to wait a while since the connection isn't available and
    48  			// on local host it is likely shutting down. It is not expected for
    49  			// this to ever reach even close to 3 seconds.
    50  		}
    51  
    52  		// It is an error we don't know how to handle, so return it
    53  		return err
    54  	}
    55  
    56  	// Context was cancelled
    57  	if errStatus := status.FromContextError(reqCtx.Err()); errStatus.Code() == codes.Canceled {
    58  		return context.Canceled
    59  	}
    60  
    61  	return err
    62  }
    63  
    64  // HandleGrpcErr is used to handle errors made to a remote gRPC plugin. It
    65  // handles detecting if the plugin has shutdown via the passeed pluginCtx. The
    66  // parameters are:
    67  // - err: the error returned from the streaming RPC
    68  // - pluginCtx: the plugins done ctx used to detect the plugin dying
    69  //
    70  // The return values are:
    71  // - ErrPluginShutdown if the error is because the plugin shutdown
    72  // - The original error
    73  func HandleGrpcErr(err error, pluginCtx context.Context) error {
    74  	if err == nil {
    75  		return nil
    76  	}
    77  
    78  	if errStatus := status.FromContextError(pluginCtx.Err()); errStatus.Code() == codes.Canceled {
    79  		// See if the plugin shutdown
    80  		select {
    81  		case <-pluginCtx.Done():
    82  			err = structs.ErrPluginShutdown
    83  		default:
    84  		}
    85  	}
    86  
    87  	// Determine if the error is because the plugin shutdown
    88  	if errStatus, ok := status.FromError(err); ok && errStatus.Code() == codes.Unavailable {
    89  		// Potentially wait a little before returning an error so we can detect
    90  		// the exit
    91  		select {
    92  		case <-pluginCtx.Done():
    93  			err = structs.ErrPluginShutdown
    94  		case <-time.After(3 * time.Second):
    95  			// Its okay to wait a while since the connection isn't available and
    96  			// on local host it is likely shutting down. It is not expected for
    97  			// this to ever reach even close to 3 seconds.
    98  		}
    99  
   100  		// It is an error we don't know how to handle, so return it
   101  		return err
   102  	}
   103  
   104  	return err
   105  }