pkg/loop/internal/core/services/capability: use NewClientConn in place of Dial for reconnect

Author: jmank88

pkg/loop/internal/core/services/capability/capabilities.go

@@ -22,7 +22,7 @@ type TriggerCapabilityClient struct {
 	*baseCapabilityClient
 }
 
-func NewTriggerCapabilityClient(brokerExt *net.BrokerExt, conn *grpc.ClientConn) capabilities.TriggerCapability {
+func NewTriggerCapabilityClient(brokerExt *net.BrokerExt, conn grpc.ClientConnInterface) capabilities.TriggerCapability {
 	return &TriggerCapabilityClient{
 		triggerExecutableClient: newTriggerExecutableClient(brokerExt, conn),
 		baseCapabilityClient:    newBaseCapabilityClient(brokerExt, conn),
@@ -39,7 +39,7 @@ type ExecutableCapability interface {
 	capabilities.BaseCapability
 }
 
-func NewExecutableCapabilityClient(brokerExt *net.BrokerExt, conn *grpc.ClientConn) ExecutableCapability {
+func NewExecutableCapabilityClient(brokerExt *net.BrokerExt, conn grpc.ClientConnInterface) ExecutableCapability {
 	return &ExecutableCapabilityClient{
 		executableClient:     newExecutableClient(brokerExt, conn),
 		baseCapabilityClient: newBaseCapabilityClient(brokerExt, conn),
@@ -52,7 +52,7 @@ type CombinedCapabilityClient struct {
 	*triggerExecutableClient
 }
 
-func NewCombinedCapabilityClient(brokerExt *net.BrokerExt, conn *grpc.ClientConn) ExecutableCapability {
+func NewCombinedCapabilityClient(brokerExt *net.BrokerExt, conn grpc.ClientConnInterface) ExecutableCapability {
 	return &CombinedCapabilityClient{
 		executableClient:        newExecutableClient(brokerExt, conn),
 		baseCapabilityClient:    newBaseCapabilityClient(brokerExt, conn),
@@ -141,7 +141,7 @@ type baseCapabilityClient struct {
 
 var _ capabilities.BaseCapability = (*baseCapabilityClient)(nil)
 
-func newBaseCapabilityClient(brokerExt *net.BrokerExt, conn *grpc.ClientConn) *baseCapabilityClient {
+func newBaseCapabilityClient(brokerExt *net.BrokerExt, conn grpc.ClientConnInterface) *baseCapabilityClient {
 	return &baseCapabilityClient{grpc: capabilitiespb.NewBaseCapabilityClient(conn), BrokerExt: brokerExt}
 }
 
@@ -348,7 +348,7 @@ func (t *triggerExecutableClient) UnregisterTrigger(ctx context.Context, req cap
 	return nil
 }
 
-func newTriggerExecutableClient(brokerExt *net.BrokerExt, conn *grpc.ClientConn) *triggerExecutableClient {
+func newTriggerExecutableClient(brokerExt *net.BrokerExt, conn grpc.ClientConnInterface) *triggerExecutableClient {
 	return &triggerExecutableClient{
 		grpc:        capabilitiespb.NewTriggerExecutableClient(conn),
 		BrokerExt:   brokerExt,
@@ -439,7 +439,7 @@ type executableClient struct {
 	*net.BrokerExt
 }
 
-func newExecutableClient(brokerExt *net.BrokerExt, conn *grpc.ClientConn) *executableClient {
+func newExecutableClient(brokerExt *net.BrokerExt, conn grpc.ClientConnInterface) *executableClient {
 	return &executableClient{
 		grpc:      capabilitiespb.NewExecutableClient(conn),
 		BrokerExt: brokerExt,

pkg/loop/internal/core/services/capability/capabilities_registry.go

@@ -3,6 +3,7 @@ package capability
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/types/known/durationpb"
@@ -208,52 +209,56 @@ func (cr *capabilitiesRegistryClient) Get(ctx context.Context, ID string) (capab
 		Id: ID,
 	}
 
-	res, err := cr.grpc.Get(ctx, req)
-	if err != nil {
-		return nil, err
-	}
-
-	conn, err := cr.Dial(res.CapabilityID)
-	if err != nil {
-		return nil, net.ErrConnDial{Name: "Capability", ID: res.CapabilityID, Err: err}
-	}
+	conn := cr.NewClientConn("Capability", func(ctx context.Context) (id uint32, deps net.Resources, err error) {
+		res, err := cr.grpc.Get(ctx, req)
+		if err != nil {
+			return 0, nil, err
+		}
+		return res.CapabilityID, nil, nil
+	})
 	client := newBaseCapabilityClient(cr.BrokerExt, conn)
-	return client, nil
+	ctx, cancel := context.WithTimeout(ctx, time.Second)
+	defer cancel()
+	_, err := client.Info(ctx) // ensure exists by triggering lazy connection with reduced timeout
+	return client, err
 }
 
 func (cr *capabilitiesRegistryClient) GetTrigger(ctx context.Context, ID string) (capabilities.TriggerCapability, error) {
 	req := &pb.GetTriggerRequest{
 		Id: ID,
 	}
 
-	res, err := cr.grpc.GetTrigger(ctx, req)
-	if err != nil {
-		return nil, err
-	}
-
-	conn, err := cr.Dial(res.CapabilityID)
-	if err != nil {
-		return nil, net.ErrConnDial{Name: "GetTrigger", ID: res.CapabilityID, Err: err}
-	}
+	conn := cr.NewClientConn("Trigger", func(ctx context.Context) (id uint32, deps net.Resources, err error) {
+		res, err := cr.grpc.GetTrigger(ctx, req)
+		if err != nil {
+			return 0, nil, err
+		}
+		return res.CapabilityID, nil, nil
+	})
 	client := NewTriggerCapabilityClient(cr.BrokerExt, conn)
-	return client, nil
+	ctx, cancel := context.WithTimeout(ctx, time.Second)
+	defer cancel()
+	_, err := client.Info(ctx) // ensure exists by triggering lazy connection with reduced timeout
+	return client, err
 }
 
 func (cr *capabilitiesRegistryClient) GetExecutable(ctx context.Context, ID string) (capabilities.ExecutableCapability, error) {
 	req := &pb.GetExecutableRequest{
 		Id: ID,
 	}
 
-	res, err := cr.grpc.GetExecutable(ctx, req)
-	if err != nil {
-		return nil, err
-	}
-	conn, err := cr.Dial(res.CapabilityID)
-	if err != nil {
-		return nil, net.ErrConnDial{Name: "GetExecutable", ID: res.CapabilityID, Err: err}
-	}
+	conn := cr.NewClientConn("Executable", func(ctx context.Context) (id uint32, deps net.Resources, err error) {
+		res, err := cr.grpc.GetExecutable(ctx, req)
+		if err != nil {
+			return 0, nil, err
+		}
+		return res.CapabilityID, nil, nil
+	})
 	client := NewExecutableCapabilityClient(cr.BrokerExt, conn)
-	return client, nil
+	ctx, cancel := context.WithTimeout(ctx, time.Second)
+	defer cancel()
+	_, err := client.Info(ctx) // ensure exists by triggering lazy connection with reduced timeout
+	return client, err
 }
 
 func (cr *capabilitiesRegistryClient) List(ctx context.Context) ([]capabilities.BaseCapability, error) {

pkg/loop/internal/net/client.go

@@ -2,6 +2,7 @@ package net
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -53,8 +54,9 @@ func (c *clientConn) Invoke(ctx context.Context, method string, args any, reply
 	cc := c.cc
 	c.mu.RUnlock()
 
+	var refErr error
 	if cc == nil {
-		cc = c.refresh(ctx, nil)
+		cc, refErr = c.refresh(ctx, nil)
 	}
 	for cc != nil {
 		err := cc.Invoke(ctx, method, args, reply, opts...)
@@ -65,41 +67,42 @@ func (c *clientConn) Invoke(ctx context.Context, method string, args any, reply
 				return err
 			}
 			c.Logger.Errorw("clientConn: Invoke: terminal error, refreshing connection", "method", method, "err", err)
-			cc = c.refresh(ctx, cc)
+			cc, refErr = c.refresh(ctx, cc)
 			continue
 		}
 		return err
 	}
-	return context.Cause(ctx)
+	return refErr
 }
 
 func (c *clientConn) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
 	c.mu.RLock()
 	cc := c.cc
 	c.mu.RUnlock()
 
+	var refErr error
 	if cc == nil {
-		cc = c.refresh(ctx, nil)
+		cc, refErr = c.refresh(ctx, nil)
 	}
 	for cc != nil {
 		s, err := cc.NewStream(ctx, desc, method, opts...)
 		if isErrTerminal(err) {
 			c.Logger.Errorw("clientConn: NewStream: terminal error, refreshing connection", "err", err)
-			cc = c.refresh(ctx, cc)
+			cc, refErr = c.refresh(ctx, cc)
 			continue
 		}
 		return s, err
 	}
-	return nil, context.Cause(ctx)
+	return nil, refErr
 }
 
 // refresh replaces c.cc with a new (different from orig) *grpc.ClientConn, and returns it as well.
 // It will block until a new connection is successfully dialed, or return nil if the context expires.
-func (c *clientConn) refresh(ctx context.Context, orig *grpc.ClientConn) *grpc.ClientConn {
+func (c *clientConn) refresh(ctx context.Context, orig *grpc.ClientConn) (*grpc.ClientConn, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	if c.cc != orig {
-		return c.cc
+		return c.cc, nil
 	}
 	if c.cc != nil {
 		if err := c.cc.Close(); err != nil {
@@ -108,13 +111,15 @@ func (c *clientConn) refresh(ctx context.Context, orig *grpc.ClientConn) *grpc.C
 		c.CloseAll(c.deps...)
 	}
 
-	try := func() bool {
-		c.Logger.Debug("Client refresh")
+	try := func() error {
+		if d, ok := ctx.Deadline(); ok {
+			c.Logger.Debugw("Client refresh", "deadline", d, "until", time.Until(d))
+		}
 		id, deps, err := c.newClient(ctx)
 		if err != nil {
 			c.Logger.Errorw("Client refresh attempt failed", "err", err)
 			c.CloseAll(deps...)
-			return false
+			return err
 		}
 		c.deps = deps
 
@@ -126,29 +131,30 @@ func (c *clientConn) refresh(ctx context.Context, orig *grpc.ClientConn) *grpc.C
 				lggr.Errorw("Client dial failed", "err", ErrConnDial{Name: c.name, ID: id, Err: err})
 			}
 			c.CloseAll(c.deps...)
-			return false
+			return err
 		}
-		return true
+		return nil
 	}
 
 	b := backoff.Backoff{
 		Min:    100 * time.Millisecond,
 		Max:    5 * time.Second,
 		Factor: 2,
 	}
-	for !try() {
+	for err := try(); err != nil; err = try() {
 		if ctx.Err() != nil {
-			c.Logger.Errorw("Client refresh failed: aborting refresh due to context error", "err", ctx.Err())
-			return nil
+			err = fmt.Errorf("%w: last error: %w", context.Cause(ctx), err)
+			c.Logger.Errorw("Client refresh failed: aborting refresh", "err", err)
+			return nil, err
 		}
 		wait := b.Duration()
 		c.Logger.Infow("Waiting to refresh", "wait", wait)
 		select {
 		case <-ctx.Done():
-			return nil
+			return nil, fmt.Errorf("%w: last error: %w", context.Cause(ctx), err)
 		case <-time.After(wait):
 		}
 	}
 
-	return c.cc
+	return c.cc, nil
 }