Merge pull request dapr#8639 from JoshVanL/fix-scheduler-client-close

Fix Scheduler Client connection pruning

Author: cicoyle

docs/release_notes/v1.15.4.md

@@ -72,7 +72,6 @@ The sidecar injector was not properly respecting the global actors enabled confi
 
 The sidecar injector now properly respects the `global.actors.enabled` helm configuration and `ACTORS_ENABLED` environment variable. When set to `false`, it will not attempt to connect to the placement service, allowing the sidecar to start successfully without actor functionality.
 
-
 ## Prevent panic of reminder operations on slow Actor Startup
 
 ### Problem

pkg/runtime/scheduler/internal/clients/clients.go

@@ -17,6 +17,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 	"sync/atomic"
 
 	schedulerv1pb "github.com/dapr/dapr/pkg/proto/scheduler/v1"
@@ -37,30 +38,34 @@ type Options struct {
 // fashion.
 type Clients struct {
 	clients     []schedulerv1pb.SchedulerClient
+	closeFns    []context.CancelFunc
 	lastUsedIdx atomic.Uint64
 }
 
 func New(ctx context.Context, opts Options) (*Clients, error) {
-	clients := make([]schedulerv1pb.SchedulerClient, len(opts.Addresses))
+	c := &Clients{
+		clients:  make([]schedulerv1pb.SchedulerClient, 0, len(opts.Addresses)),
+		closeFns: make([]context.CancelFunc, 0, len(opts.Addresses)),
+	}
 
-	for i, address := range opts.Addresses {
+	for _, address := range opts.Addresses {
 		log.Debugf("Attempting to connect to Scheduler at address: %s", address)
-		client, err := client.New(ctx, address, opts.Security)
+		client, closeFn, err := client.New(ctx, address, opts.Security)
 		if err != nil {
+			c.Close()
 			return nil, fmt.Errorf("scheduler client not initialized for address %s: %s", address, err)
 		}
 
 		log.Infof("Scheduler client initialized for address: %s", address)
-		clients[i] = client
+		c.clients = append(c.clients, client)
+		c.closeFns = append(c.closeFns, closeFn)
 	}
 
-	if len(clients) > 0 {
+	if len(c.clients) > 0 {
 		log.Info("Scheduler clients initialized")
 	}
 
-	return &Clients{
-		clients: clients,
-	}, nil
+	return c, nil
 }
 
 // Next returns the next client in a round-robin manner.
@@ -77,3 +82,15 @@ func (c *Clients) Next() (schedulerv1pb.SchedulerClient, error) {
 func (c *Clients) All() []schedulerv1pb.SchedulerClient {
 	return c.clients
 }
+
+func (c *Clients) Close() {
+	var wg sync.WaitGroup
+	wg.Add(len(c.closeFns))
+	for _, closeFn := range c.closeFns {
+		go func() {
+			closeFn()
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+}

pkg/runtime/scheduler/internal/cluster/cluster.go

@@ -133,7 +133,7 @@ func (c *Cluster) RunClients(ctx context.Context, clients *clients.Clients) erro
 		}
 
 		if err != nil {
-			log.Warnf("Error watching scheduler jobs: %v", err)
+			return fmt.Errorf("error watching scheduler jobs: %w", err)
 		}
 	}
 }

pkg/runtime/scheduler/internal/cluster/streamer.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"io"
 	"sync"
+	"sync/atomic"
 
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -41,7 +42,8 @@ type streamer struct {
 	channels *channels.Channels
 	wfengine wfengine.Interface
 
-	wg sync.WaitGroup
+	wg       sync.WaitGroup
+	inflight atomic.Int64
 }
 
 // run starts the streamer and blocks until the stream is closed or an error occurs.
@@ -53,7 +55,10 @@ func (s *streamer) run(ctx context.Context) error {
 // scheduler job messages. It then invokes the appropriate app or actor
 // reminder based on the job metadata.
 func (s *streamer) receive(ctx context.Context) error {
-	defer s.wg.Wait()
+	defer func() {
+		s.wg.Wait()
+		s.stream.CloseSend()
+	}()
 
 	for {
 		resp, err := s.stream.Recv()
@@ -65,12 +70,15 @@ func (s *streamer) receive(ctx context.Context) error {
 		}
 
 		s.wg.Add(1)
+		s.inflight.Add(1)
 		go func() {
-			defer s.wg.Done()
+			defer func() {
+				s.wg.Done()
+				s.inflight.Add(-1)
+			}()
+
 			result := s.handleJob(ctx, resp)
 			select {
-			case <-ctx.Done():
-			case <-s.stream.Context().Done():
 			case s.resultCh <- &schedulerv1pb.WatchJobsRequest{
 				WatchJobRequestType: &schedulerv1pb.WatchJobsRequest_Result{
 					Result: &schedulerv1pb.WatchJobsRequestResult{
@@ -79,6 +87,8 @@ func (s *streamer) receive(ctx context.Context) error {
 					},
 				},
 			}:
+			case <-s.stream.Context().Done():
+			case <-ctx.Done():
 			}
 		}()
 	}
@@ -93,14 +103,14 @@ func (s *streamer) outgoing(ctx context.Context) error {
 
 	for {
 		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-s.stream.Context().Done():
-			return s.stream.Context().Err()
 		case result := <-s.resultCh:
 			if err := s.stream.Send(result); err != nil {
 				return err
 			}
+		case <-ctx.Done():
+			if s.inflight.Load() == 0 && len(s.resultCh) == 0 {
+				return ctx.Err()
+			}
 		}
 	}
 }

pkg/runtime/scheduler/internal/watchhosts/watchhosts.go

@@ -0,0 +1,197 @@
+/*
+Copyright 2024 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 watchhosts
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"slices"
+	"sync"
+	"sync/atomic"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	schedulerv1pb "github.com/dapr/dapr/pkg/proto/scheduler/v1"
+	"github.com/dapr/dapr/pkg/scheduler/client"
+	"github.com/dapr/dapr/pkg/security"
+	"github.com/dapr/kit/events/broadcaster"
+	"github.com/dapr/kit/logger"
+)
+
+var log = logger.NewLogger("dapr.runtime.scheduler.watchhosts")
+
+type Options struct {
+	Addresses []string
+	Security  security.Handler
+}
+
+type WatchHosts struct {
+	allAddrs []string
+	security security.Handler
+
+	gotAddrs atomic.Pointer[ContextAddress]
+	subs     *broadcaster.Broadcaster[*ContextAddress]
+
+	cancel  context.CancelFunc
+	readyCh chan struct{}
+	lock    sync.Mutex
+}
+
+type ContextAddress struct {
+	context.Context
+	Addresses []string
+}
+
+func New(opts Options) *WatchHosts {
+	return &WatchHosts{
+		allAddrs: opts.Addresses,
+		security: opts.Security,
+		subs:     broadcaster.New[*ContextAddress](),
+		readyCh:  make(chan struct{}),
+	}
+}
+
+func (w *WatchHosts) Run(ctx context.Context) error {
+	defer w.subs.Close()
+
+	stream, closeCon, err := w.connSchedulerHosts(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to connect to scheduler host: %s", err)
+	}
+
+	if stream != nil {
+		defer func() {
+			stream.CloseSend()
+			closeCon()
+		}()
+	}
+
+	err = w.handleStream(ctx, stream)
+
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
+	if err != nil {
+		return fmt.Errorf("failed to handle Scheduler WatchHosts stream: %s", err)
+	}
+
+	return nil
+}
+
+func (w *WatchHosts) Addresses(ctx context.Context) <-chan *ContextAddress {
+	ch := make(chan *ContextAddress, 1)
+	select {
+	case <-ctx.Done():
+		ch <- &ContextAddress{ctx, nil}
+		return ch
+	case <-w.readyCh:
+	}
+
+	w.lock.Lock()
+	defer w.lock.Unlock()
+	w.subs.Subscribe(ctx, ch)
+	ch <- w.gotAddrs.Load()
+
+	return ch
+}
+
+func (w *WatchHosts) handleStream(ctx context.Context, stream schedulerv1pb.Scheduler_WatchHostsClient) error {
+	// If no stream was made the server doesn't support watching hosts
+	// (pre-1.15), so we use static. Remove in 1.16.
+	if stream == nil {
+		w.lock.Lock()
+		got := &ContextAddress{ctx, w.allAddrs}
+		w.gotAddrs.Store(got)
+		w.subs.Broadcast(got)
+		w.lock.Unlock()
+		close(w.readyCh)
+		<-ctx.Done()
+		return nil
+	}
+
+	for {
+		gotAddrs, err := w.watchNextAddresses(ctx, stream)
+		if err != nil {
+			return err
+		}
+
+		w.lock.Lock()
+		if w.cancel != nil {
+			w.cancel()
+		}
+
+		actx, cancel := context.WithCancel(ctx)
+		w.cancel = cancel
+		got := &ContextAddress{actx, gotAddrs}
+		w.gotAddrs.Store(got)
+		w.subs.Broadcast(got)
+		w.lock.Unlock()
+
+		select {
+		case <-w.readyCh:
+		default:
+			close(w.readyCh)
+		}
+	}
+}
+
+func (w *WatchHosts) watchNextAddresses(ctx context.Context, stream schedulerv1pb.Scheduler_WatchHostsClient) ([]string, error) {
+	resp, err := stream.Recv()
+	if err != nil {
+		if status.Code(err) == codes.Unimplemented {
+			// Ignore unimplemented error code as we are talking to an old server.
+			// TODO: @joshvanl: remove special case in v1.16.
+			return slices.Clone(w.allAddrs), nil
+		}
+		return nil, err
+	}
+
+	gotAddrs := make([]string, 0, len(resp.GetHosts()))
+	for _, host := range resp.GetHosts() {
+		gotAddrs = append(gotAddrs, host.GetAddress())
+	}
+
+	log.Infof("Received updated scheduler hosts addresses: %v", gotAddrs)
+
+	return gotAddrs, nil
+}
+
+func (w *WatchHosts) connSchedulerHosts(ctx context.Context) (schedulerv1pb.Scheduler_WatchHostsClient, context.CancelFunc, error) {
+	//nolint:gosec
+	i := rand.Intn(len(w.allAddrs))
+	log.Debugf("Attempting to connect to scheduler to WatchHosts: %s", w.allAddrs[i])
+
+	// This is connecting to a DNS A rec which will return healthy scheduler
+	// hosts.
+	cl, closeCon, err := client.New(ctx, w.allAddrs[i], w.security)
+	if err != nil {
+		return nil, nil, fmt.Errorf("scheduler client not initialized for address %s: %s", w.allAddrs[i], err)
+	}
+
+	stream, err := cl.WatchHosts(ctx, new(schedulerv1pb.WatchHostsRequest))
+	if err != nil {
+		if status.Code(err) == codes.Unimplemented {
+			// Ignore unimplemented error code as we are talking to an old server.
+			// TODO: @joshvanl: remove special case in v1.16.
+			return nil, nil, nil
+		}
+
+		return nil, nil, fmt.Errorf("failed to watch scheduler hosts: %s", err)
+	}
+
+	return stream, closeCon, nil
+}