github.com/hernad/nomad@v1.6.112/helper/pluginutils/grpcutils/utils.go (about)

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