Plugin client's handle plugin dying

drivers/exec/driver.go

@@ -1,6 +1,7 @@
 package exec
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -20,7 +21,6 @@ import (
 	"github.com/hashicorp/nomad/plugins/shared"
 	"github.com/hashicorp/nomad/plugins/shared/hclspec"
 	"github.com/hashicorp/nomad/plugins/shared/loader"
-	"golang.org/x/net/context"
 )
 
 const (

drivers/java/driver.go

@@ -1,6 +1,7 @@
 package java
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
@@ -23,7 +24,6 @@ import (
 	"github.com/hashicorp/nomad/plugins/shared"
 	"github.com/hashicorp/nomad/plugins/shared/hclspec"
 	"github.com/hashicorp/nomad/plugins/shared/loader"
-	"golang.org/x/net/context"
 )
 
 const (

drivers/mock/driver.go

@@ -16,7 +16,6 @@ import (
 	"github.com/hashicorp/nomad/plugins/drivers"
 	"github.com/hashicorp/nomad/plugins/shared/hclspec"
 	"github.com/hashicorp/nomad/plugins/shared/loader"
-	netctx "golang.org/x/net/context"
 )
 
 const (
@@ -232,7 +231,7 @@ func (d *Driver) Capabilities() (*drivers.Capabilities, error) {
 	return capabilities, nil
 }
 
-func (d *Driver) Fingerprint(ctx netctx.Context) (<-chan *drivers.Fingerprint, error) {
+func (d *Driver) Fingerprint(ctx context.Context) (<-chan *drivers.Fingerprint, error) {
 	ch := make(chan *drivers.Fingerprint)
 	go d.handleFingerprint(ctx, ch)
 	return ch, nil
@@ -365,7 +364,7 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *cstru
 
 }
 
-func (d *Driver) WaitTask(ctx netctx.Context, taskID string) (<-chan *drivers.ExitResult, error) {
+func (d *Driver) WaitTask(ctx context.Context, taskID string) (<-chan *drivers.ExitResult, error) {
 	handle, ok := d.tasks.Get(taskID)
 	if !ok {
 		return nil, drivers.ErrTaskNotFound
@@ -430,7 +429,7 @@ func (d *Driver) TaskStats(taskID string) (*cstructs.TaskResourceUsage, error) {
 	return nil, nil
 }
 
-func (d *Driver) TaskEvents(ctx netctx.Context) (<-chan *drivers.TaskEvent, error) {
+func (d *Driver) TaskEvents(ctx context.Context) (<-chan *drivers.TaskEvent, error) {
 	return d.eventer.TaskEvents(ctx)
 }
 

drivers/qemu/driver.go

@@ -1,6 +1,7 @@
 package qemu
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
@@ -25,7 +26,6 @@ import (
 	"github.com/hashicorp/nomad/plugins/shared"
 	"github.com/hashicorp/nomad/plugins/shared/hclspec"
 	"github.com/hashicorp/nomad/plugins/shared/loader"
-	"golang.org/x/net/context"
 )
 
 const (

drivers/rawexec/driver.go

@@ -1,6 +1,7 @@
 package rawexec
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -22,7 +23,6 @@ import (
 	"github.com/hashicorp/nomad/plugins/shared"
 	"github.com/hashicorp/nomad/plugins/shared/hclspec"
 	"github.com/hashicorp/nomad/plugins/shared/loader"
-	"golang.org/x/net/context"
 )
 
 const (

drivers/rkt/driver.go

@@ -4,6 +4,7 @@ package rkt
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -36,7 +37,6 @@ import (
 	"github.com/hashicorp/nomad/plugins/shared/hclspec"
 	"github.com/hashicorp/nomad/plugins/shared/loader"
 	rktv1 "github.com/rkt/rkt/api/v1"
-	"golang.org/x/net/context"
 )
 
 const (

drivers/rkt/driver_test.go

@@ -3,17 +3,16 @@
 package rkt
 
 import (
+	"bytes"
+	"context"
 	"fmt"
 	"io/ioutil"
+	"os"
 	"path/filepath"
 	"sync"
 	"testing"
 	"time"
 
-	"os"
-
-	"bytes"
-
 	"github.com/hashicorp/hcl2/hcl"
 	ctestutil "github.com/hashicorp/nomad/client/testutil"
 	"github.com/hashicorp/nomad/helper/testlog"
@@ -26,7 +25,6 @@ import (
 	"github.com/hashicorp/nomad/plugins/shared/hclspec"
 	"github.com/hashicorp/nomad/testutil"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/net/context"
 )
 
 var _ drivers.DriverPlugin = (*Driver)(nil)

drivers/shared/eventer/eventer.go

@@ -1,12 +1,12 @@
 package eventer
 
 import (
+	"context"
 	"sync"
 	"time"
 
 	hclog "github.com/hashicorp/go-hclog"
 	"github.com/hashicorp/nomad/plugins/drivers"
-	"golang.org/x/net/context"
 )
 
 var (

plugins/base/client.go

@@ -12,10 +12,13 @@ import (
 // gRPC to communicate to the remote plugin.
 type BasePluginClient struct {
 	Client proto.BasePluginClient
+
+	// DoneCtx is closed when the plugin exits
+	DoneCtx context.Context
 }
 
 func (b *BasePluginClient) PluginInfo() (*PluginInfoResponse, error) {
-	presp, err := b.Client.PluginInfo(context.Background(), &proto.PluginInfoRequest{})
+	presp, err := b.Client.PluginInfo(b.DoneCtx, &proto.PluginInfoRequest{})
 	if err != nil {
 		return nil, err
 	}
@@ -41,7 +44,7 @@ func (b *BasePluginClient) PluginInfo() (*PluginInfoResponse, error) {
 }
 
 func (b *BasePluginClient) ConfigSchema() (*hclspec.Spec, error) {
-	presp, err := b.Client.ConfigSchema(context.Background(), &proto.ConfigSchemaRequest{})
+	presp, err := b.Client.ConfigSchema(b.DoneCtx, &proto.ConfigSchemaRequest{})
 	if err != nil {
 		return nil, err
 	}
@@ -51,7 +54,7 @@ func (b *BasePluginClient) ConfigSchema() (*hclspec.Spec, error) {
 
 func (b *BasePluginClient) SetConfig(data []byte, config *ClientAgentConfig) error {
 	// Send the config
-	_, err := b.Client.SetConfig(context.Background(), &proto.SetConfigRequest{
+	_, err := b.Client.SetConfig(b.DoneCtx, &proto.SetConfigRequest{
 		MsgpackConfig: data,
 		NomadConfig:   config.toProto(),
 	})

plugins/base/plugin.go

@@ -51,7 +51,10 @@ func (p *PluginBase) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Server) error
 }
 
 func (p *PluginBase) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker, c *grpc.ClientConn) (interface{}, error) {
-	return &BasePluginClient{Client: proto.NewBasePluginClient(c)}, nil
+	return &BasePluginClient{
+		Client:  proto.NewBasePluginClient(c),
+		DoneCtx: ctx,
+	}, nil
 }
 
 // MsgpackHandle is a shared handle for encoding/decoding of structs

plugins/device/client.go

@@ -9,9 +9,7 @@ import (
 	"github.com/golang/protobuf/ptypes"
 	"github.com/hashicorp/nomad/plugins/base"
 	"github.com/hashicorp/nomad/plugins/device/proto"
-	netctx "golang.org/x/net/context"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
+	"github.com/hashicorp/nomad/plugins/shared"
 )
 
 // devicePluginClient implements the client side of a remote device plugin, using
@@ -49,28 +47,33 @@ func (d *devicePluginClient) Fingerprint(ctx context.Context) (<-chan *Fingerpri
 // the gRPC stream to a channel. Exits either when context is cancelled or the
 // stream has an error.
 func (d *devicePluginClient) handleFingerprint(
-	ctx netctx.Context,
+	ctx context.Context,
 	stream proto.DevicePlugin_FingerprintClient,
 	out chan *FingerprintResponse) {
 
+	defer close(out)
 	for {
 		resp, err := stream.Recv()
 		if err != nil {
 			if err != io.EOF {
 				out <- &FingerprintResponse{
-					Error: d.handleStreamErr(err, ctx),
+					Error: shared.HandleStreamErr(err, ctx, d.doneCtx),
 				}
 			}
 
 			// End the stream
-			close(out)
 			return
 		}
 
 		// Send the response
-		out <- &FingerprintResponse{
+		f := &FingerprintResponse{
 			Devices: convertProtoDeviceGroups(resp.GetDeviceGroup()),
 		}
+		select {
+		case <-ctx.Done():
+			return
+		case out <- f:
+		}
 	}
 }
 
@@ -116,69 +119,32 @@ func (d *devicePluginClient) Stats(ctx context.Context, interval time.Duration)
 // the gRPC stream to a channel. Exits either when context is cancelled or the
 // stream has an error.
 func (d *devicePluginClient) handleStats(
-	ctx netctx.Context,
+	ctx context.Context,
 	stream proto.DevicePlugin_StatsClient,
 	out chan *StatsResponse) {
 
+	defer close(out)
 	for {
 		resp, err := stream.Recv()
 		if err != nil {
 			if err != io.EOF {
 				out <- &StatsResponse{
-					Error: d.handleStreamErr(err, ctx),
+					Error: shared.HandleStreamErr(err, ctx, d.doneCtx),
 				}
 			}
 
 			// End the stream
-			close(out)
 			return
 		}
 
 		// Send the response
-		out <- &StatsResponse{
+		s := &StatsResponse{
 			Groups: convertProtoDeviceGroupsStats(resp.GetGroups()),
 		}
-	}
-}
-
-// handleStreamErr is used to handle a non io.EOF error in a stream. It handles
-// detecting if the plugin has shutdown
-func (d *devicePluginClient) handleStreamErr(err error, ctx context.Context) error {
-	if err == nil {
-		return nil
-	}
-
-	// Determine if the error is because the plugin shutdown
-	if errStatus, ok := status.FromError(err); ok && errStatus.Code() == codes.Unavailable {
-		// Potentially wait a little before returning an error so we can detect
-		// the exit
 		select {
-		case <-d.doneCtx.Done():
-			err = base.ErrPluginShutdown
 		case <-ctx.Done():
-			err = ctx.Err()
-
-			// There is no guarantee that the select will choose the
-			// doneCtx first so we have to double check
-			select {
-			case <-d.doneCtx.Done():
-				err = base.ErrPluginShutdown
-			default:
-			}
-		case <-time.After(3 * time.Second):
-			// Its okay to wait a while since the connection isn't available and
-			// on local host it is likely shutting down. It is not expected for
-			// this to ever reach even close to 3 seconds.
+			return
+		case out <- s:
 		}
-
-		// It is an error we don't know how to handle, so return it
-		return err
 	}
-
-	// Context was cancelled
-	if errStatus := status.FromContextError(ctx.Err()); errStatus.Code() == codes.Canceled {
-		return context.Canceled
-	}
-
-	return err
 }

plugins/device/plugin.go

@@ -31,7 +31,8 @@ func (p *PluginDevice) GRPCClient(ctx context.Context, broker *plugin.GRPCBroker
 		doneCtx: ctx,
 		client:  proto.NewDevicePluginClient(c),
 		BasePluginClient: &base.BasePluginClient{
-			Client: bproto.NewBasePluginClient(c),
+			Client:  bproto.NewBasePluginClient(c),
+			DoneCtx: ctx,
 		},
 	}, nil
 }