Graceful shutdown: don't cancel in-flight input binding events (dapr#8637)

* Graceful shutdown: don't cancel in-flight input binding events

Do not cancel in-flight input binding requests during graceful shutdown
procedure, instead wait for them to complete or until the grace period
expires. Also ensures in-flight requests are resolved during hot-reload
events. New requests while the shutdown is in progress are errors to the
input binding Component.

Fixes: dapr#8332

Signed-off-by: joshvanl <[email protected]>

* Adds subscription len check on block

Signed-off-by: joshvanl <[email protected]>

---------

Signed-off-by: joshvanl <[email protected]>

Author: JoshVanL

pkg/components/bindings/input_pluggable.go

@@ -29,7 +29,6 @@ import (
 // grpcInputBinding is a implementation of a inputbinding over a gRPC Protocol.
 type grpcInputBinding struct {
 	*pluggable.GRPCConnector[proto.InputBindingClient]
-	bindings.InputBinding
 	logger logger.Logger
 
 	closed  atomic.Bool
@@ -130,7 +129,7 @@ func (b *grpcInputBinding) Read(ctx context.Context, handler bindings.Handler) e
 
 			// TODO reconnect on error
 			if err != nil {
-				b.logger.Errorf("failed to receive message: %v", err)
+				b.logger.Errorf("failed to receive binding message: %v", err)
 				return
 			}
 			b.wg.Add(1)
@@ -149,7 +148,7 @@ func (b *grpcInputBinding) Close() error {
 	if b.closed.CompareAndSwap(false, true) {
 		close(b.closeCh)
 	}
-	return b.InputBinding.Close()
+	return nil
 }
 
 // inputFromConnector creates a new GRPC inputbinding using the given underlying connector.

pkg/components/pubsub/pluggable.go

@@ -159,10 +159,8 @@ func (p *grpcPubSub) adaptHandler(ctx context.Context, streamingPull proto.PubSu
 func (p *grpcPubSub) pullMessages(parentCtx context.Context, topic *proto.Topic, handler pubsub.Handler) error {
 	ctx, cancel := context.WithCancel(context.Background())
 
-	var wg sync.WaitGroup
 	go func() {
 		<-parentCtx.Done()
-		wg.Wait()
 		cancel()
 	}()
 
@@ -207,11 +205,7 @@ func (p *grpcPubSub) pullMessages(parentCtx context.Context, topic *proto.Topic,
 
 			p.logger.Debugf("Received message from stream on topic %s", msg.GetTopicName())
 
-			wg.Add(1)
-			go func() {
-				handle(msg)
-				wg.Done()
-			}()
+			handle(msg)
 		}
 	}()
 

pkg/runtime/processor/binding/binding.go

@@ -32,6 +32,7 @@ import (
 	"github.com/dapr/dapr/pkg/runtime/compstore"
 	rterrors "github.com/dapr/dapr/pkg/runtime/errors"
 	"github.com/dapr/dapr/pkg/runtime/meta"
+	"github.com/dapr/dapr/pkg/runtime/processor/binding/input"
 	"github.com/dapr/kit/logger"
 )
 
@@ -75,7 +76,7 @@ type binding struct {
 	stopForever     bool
 
 	subscribeBindingList []string
-	inputCancels         map[string]context.CancelFunc
+	activeInputs         map[string]*input.Input
 	wg                   sync.WaitGroup
 }
 
@@ -89,7 +90,7 @@ func New(opts Options) *binding {
 		tracingSpec:  opts.TracingSpec,
 		grpc:         opts.GRPC,
 		channels:     opts.Channels,
-		inputCancels: make(map[string]context.CancelFunc),
+		activeInputs: make(map[string]*input.Input),
 	}
 }
 
@@ -130,10 +131,10 @@ func (b *binding) Close(comp compapi.Component) error {
 	inbinding, ok := b.compStore.GetInputBinding(comp.Name)
 	if ok {
 		defer b.compStore.DeleteInputBinding(comp.Name)
-		if cancel := b.inputCancels[comp.Name]; cancel != nil {
-			cancel()
+		if input := b.activeInputs[comp.Name]; input != nil {
+			input.Stop()
 		}
-		delete(b.inputCancels, comp.Name)
+		delete(b.activeInputs, comp.Name)
 		if err := inbinding.Close(); err != nil {
 			errs = append(errs, err)
 		}

pkg/runtime/processor/binding/input/input.go

@@ -0,0 +1,105 @@
+/*
+Copyright 2025 The Dapr Authors
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package input
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/dapr/components-contrib/bindings"
+	diag "github.com/dapr/dapr/pkg/diagnostics"
+	"github.com/dapr/kit/logger"
+)
+
+var log = logger.NewLogger("dapr.runtime.processor.binding.input")
+
+type Options struct {
+	Name    string
+	Binding bindings.InputBinding
+	Handler func(context.Context, string, []byte, map[string]string) ([]byte, error)
+}
+
+type Input struct {
+	name    string
+	binding bindings.InputBinding
+	handler func(context.Context, string, []byte, map[string]string) ([]byte, error)
+
+	cancel   func()
+	closed   atomic.Bool
+	wg       sync.WaitGroup
+	inflight atomic.Int64
+}
+
+func Run(opts Options) (*Input, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	i := &Input{
+		name:    opts.Name,
+		binding: opts.Binding,
+		handler: opts.Handler,
+		cancel:  cancel,
+	}
+
+	return i, i.read(ctx)
+}
+
+func (i *Input) Stop() {
+	i.closed.Store(true)
+	inflight := i.inflight.Load() > 0
+
+	i.wg.Wait()
+
+	// If there were in-flight requests then wait some time for the result to be
+	// sent to the binding. This is because the message result context is
+	// disparate.
+	if inflight {
+		time.Sleep(time.Millisecond * 400)
+	}
+	i.cancel()
+}
+
+func (i *Input) read(ctx context.Context) error {
+	return i.binding.Read(ctx, func(ctx context.Context, resp *bindings.ReadResponse) ([]byte, error) {
+		i.wg.Add(1)
+		i.inflight.Add(1)
+		defer func() {
+			i.wg.Done()
+			i.inflight.Add(-1)
+		}()
+
+		if i.closed.Load() {
+			return nil, errors.New("input binding is closed")
+		}
+
+		if resp == nil {
+			return nil, nil
+		}
+
+		start := time.Now()
+		b, err := i.handler(ctx, i.name, resp.Data, resp.Metadata)
+		elapsed := diag.ElapsedSince(start)
+
+		diag.DefaultComponentMonitoring.InputBindingEvent(context.Background(), i.name, err == nil, elapsed)
+
+		if err != nil {
+			log.Debugf("error from app consumer for binding [%s]: %s", i.name, err)
+			return nil, err
+		}
+
+		return b, nil
+	})
+}

pkg/runtime/processor/binding/send.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"go.opentelemetry.io/otel/trace"
@@ -37,6 +38,7 @@ import (
 	invokev1 "github.com/dapr/dapr/pkg/messaging/v1"
 	runtimev1pb "github.com/dapr/dapr/pkg/proto/runtime/v1"
 	"github.com/dapr/dapr/pkg/resiliency"
+	"github.com/dapr/dapr/pkg/runtime/processor/binding/input"
 )
 
 func (b *binding) StartReadingFromBindings(ctx context.Context) error {
@@ -54,10 +56,16 @@ func (b *binding) StartReadingFromBindings(ctx context.Context) error {
 	}
 
 	// Clean any previous state
-	for _, cancel := range b.inputCancels {
-		cancel()
+	var wg sync.WaitGroup
+	wg.Add(len(b.activeInputs))
+	for _, inp := range b.activeInputs {
+		go func(input *input.Input) {
+			input.Stop()
+			wg.Done()
+		}(inp)
 	}
-	b.inputCancels = make(map[string]context.CancelFunc)
+	wg.Wait()
+	clear(b.activeInputs)
 
 	comps := b.compStore.ListComponents()
 	bindings := make(map[string]componentsV1alpha1.Component)
@@ -86,31 +94,34 @@ func (b *binding) startInputBinding(comp componentsV1alpha1.Component, binding b
 
 	m := meta.Properties
 
-	ctx, cancel := context.WithCancel(context.Background())
 	if isBindingOfExplicitDirection(ComponentTypeInput, m) {
 		isSubscribed = true
 	} else {
-		var err error
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
+		defer cancel()
 		isSubscribed, err = b.isAppSubscribedToBinding(ctx, comp.Name)
 		if err != nil {
-			cancel()
 			return err
 		}
 	}
 
 	if !isSubscribed {
 		log.Infof("app has not subscribed to binding %s.", comp.Name)
-		cancel()
 		return nil
 	}
 
-	if err := b.readFromBinding(ctx, comp.Name, binding); err != nil {
+	input, err := input.Run(input.Options{
+		Name:    comp.Name,
+		Binding: binding,
+		Handler: b.sendBindingEventToApp,
+	})
+	if err != nil {
 		log.Errorf("error reading from input binding %s: %s", comp.Name, err)
-		cancel()
-		return nil
+		return err
 	}
 
-	b.inputCancels[comp.Name] = cancel
+	b.activeInputs[comp.Name] = input
+
 	return nil
 }
 
@@ -125,10 +136,16 @@ func (b *binding) StopReadingFromBindings(forever bool) {
 
 	b.readingBindings = false
 
-	for _, cancel := range b.inputCancels {
-		cancel()
+	var wg sync.WaitGroup
+	wg.Add(len(b.activeInputs))
+	for _, inp := range b.activeInputs {
+		go func(input *input.Input) {
+			input.Stop()
+			wg.Done()
+		}(inp)
 	}
-	b.inputCancels = make(map[string]context.CancelFunc)
+	wg.Wait()
+	clear(b.activeInputs)
 }
 
 func (b *binding) sendBatchOutputBindingsParallel(ctx context.Context, to []string, data []byte) {
@@ -406,26 +423,6 @@ func (b *binding) sendBindingEventToApp(ctx context.Context, bindingName string,
 	return appResponseBody, nil
 }
 
-func (b *binding) readFromBinding(readCtx context.Context, name string, binding bindings.InputBinding) error {
-	return binding.Read(readCtx, func(ctx context.Context, resp *bindings.ReadResponse) ([]byte, error) {
-		if resp == nil {
-			return nil, nil
-		}
-
-		start := time.Now()
-		b, err := b.sendBindingEventToApp(ctx, name, resp.Data, resp.Metadata)
-		elapsed := diag.ElapsedSince(start)
-
-		diag.DefaultComponentMonitoring.InputBindingEvent(context.Background(), name, err == nil, elapsed)
-
-		if err != nil {
-			log.Debugf("error from app consumer for binding [%s]: %s", name, err)
-			return nil, err
-		}
-		return b, nil
-	})
-}
-
 func (b *binding) getSubscribedBindingsGRPC(ctx context.Context) ([]string, error) {
 	conn, err := b.grpc.GetAppClient()
 	if err != nil {