notify/bulk_recovery.go

package notify

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

const (
	bulkRecoveryQueueSize = 64
	bulkRecoveryWorkers   = 4
	bulkRecoveryAttempts  = 3
	bulkRecoveryAttempt   = 750 * time.Millisecond
	bulkRecoveryBackoff   = 50 * time.Millisecond
)

var errBulkRecoveryQueueFull = errors.New("bulk recovery queue full")

type bulkRecoveryTask func(context.Context) error

// bulkRecoveryQueue bounds cleanup work after failed bulk opens. A small fixed
// worker set prevents both unbounded goroutines and multi-minute serial drain.
type bulkRecoveryQueue struct {
	mu             sync.Mutex
	tasks          []bulkRecoveryTask
	workerRunning  bool
	workersRunning int
	onError        func(error)
}

func newBulkRecoveryQueue(onError func(error)) *bulkRecoveryQueue {
	return &bulkRecoveryQueue{
		onError: onError,
	}
}

func (q *bulkRecoveryQueue) enqueue(task bulkRecoveryTask) bool {
	if task == nil {
		return true
	}
	if q == nil {
		return false
	}
	q.mu.Lock()
	if len(q.tasks) >= bulkRecoveryQueueSize {
		q.mu.Unlock()
		return false
	}
	q.tasks = append(q.tasks, task)
	start := 0
	for q.workersRunning < bulkRecoveryWorkers && q.workersRunning < len(q.tasks) {
		q.workersRunning++
		start++
	}
	q.workerRunning = q.workersRunning > 0
	q.mu.Unlock()
	for i := 0; i < start; i++ {
		go q.loop()
	}
	return true
}

func (q *bulkRecoveryQueue) loop() {
	for {
		q.mu.Lock()
		if len(q.tasks) == 0 {
			q.workersRunning--
			q.workerRunning = q.workersRunning > 0
			q.mu.Unlock()
			return
		}
		task := q.tasks[0]
		copy(q.tasks, q.tasks[1:])
		q.tasks[len(q.tasks)-1] = nil
		q.tasks = q.tasks[:len(q.tasks)-1]
		q.mu.Unlock()
		q.execute(task)
	}
}

func (q *bulkRecoveryQueue) execute(task bulkRecoveryTask) {
	if q == nil || task == nil {
		return
	}
	q.run(task)
}

func (q *bulkRecoveryQueue) run(task bulkRecoveryTask) {
	if q == nil || task == nil {
		return
	}
	lastErr := runBulkRecoveryTaskContext(context.Background(), task)
	if lastErr != nil && q.onError != nil {
		q.onError(lastErr)
	}
}

// runBulkRecoveryTask performs a bounded reset attempt synchronously. Callers
// that must establish ordering (for example, Auto dedicated -> shared fallback)
// use this path so a later open cannot overtake cleanup queued in the background.
func runBulkRecoveryTask(task bulkRecoveryTask) error {
	return runBulkRecoveryTaskContext(context.Background(), task)
}

func runBulkRecoveryTaskContext(parent context.Context, task bulkRecoveryTask) error {
	if task == nil {
		return nil
	}
	if parent == nil {
		parent = context.Background()
	}
	deadline := time.Now().Add(bulkOpenRecoveryTimeout)
	var lastErr error
	for attempt := 0; attempt < bulkRecoveryAttempts; attempt++ {
		if err := parent.Err(); err != nil {
			return err
		}
		remaining := time.Until(deadline)
		if remaining <= 0 {
			break
		}
		attemptTimeout := remaining
		if attemptTimeout > bulkRecoveryAttempt {
			attemptTimeout = bulkRecoveryAttempt
		}
		ctx, cancel := context.WithTimeout(parent, attemptTimeout)
		err := task(ctx)
		cancel()
		if err == nil {
			return nil
		}
		lastErr = err
		if !bulkRecoveryErrorRetryable(err) {
			break
		}
		if attempt+1 >= bulkRecoveryAttempts {
			break
		}
		backoff := bulkRecoveryBackoff << attempt
		if backoff > time.Until(deadline) {
			backoff = time.Until(deadline)
		}
		if backoff > 0 {
			timer := time.NewTimer(backoff)
			select {
			case <-parent.Done():
				if !timer.Stop() {
					<-timer.C
				}
				return parent.Err()
			case <-timer.C:
			}
		}
	}
	return lastErr
}

func bulkRecoveryErrorRetryable(err error) bool {
	if err == nil {
		return false
	}
	return !errors.Is(err, errTransportDetached) &&
		!errors.Is(err, errServiceShutdown) &&
		!errors.Is(err, net.ErrClosed) &&
		!errors.Is(err, io.ErrClosedPipe)
}

func (c *ClientCommon) resetBulkAtRouteAndWait(ctx context.Context, route clientSessionRoute, req BulkResetRequest) error {
	if c == nil {
		return errBulkClientNil
	}
	err := runBulkRecoveryTaskContext(ctx, newClientBulkResetRecoveryTaskAtRoute(c, route, req))
	if err != nil && bulkRecoveryErrorRetryable(err) {
		c.bestEffortBulkResetAtRoute(route, req)
	}
	return err
}

func (c *ClientCommon) cleanupBulkResetAtRoute(ctx context.Context, route clientSessionRoute, req BulkResetRequest, wait bool) error {
	if wait {
		return c.resetBulkAtRouteAndWait(ctx, route, req)
	}
	c.bestEffortBulkResetAtRoute(route, req)
	return nil
}

func (s *ServerCommon) resetBulkLogicalAndWait(ctx context.Context, logical *LogicalConn, transport *TransportConn, req BulkResetRequest) error {
	if s == nil {
		return errBulkServerNil
	}
	err := runBulkRecoveryTaskContext(ctx, newServerBulkResetRecoveryTask(s, logical, transport, req))
	if err != nil && bulkRecoveryErrorRetryable(err) {
		s.bestEffortBulkResetLogical(logical, transport, req)
	}
	return err
}

func (s *ServerCommon) cleanupBulkLogicalReset(ctx context.Context, logical *LogicalConn, transport *TransportConn, req BulkResetRequest, wait bool) error {
	if wait {
		return s.resetBulkLogicalAndWait(ctx, logical, transport, req)
	}
	s.bestEffortBulkResetLogical(logical, transport, req)
	return nil
}

func (s *ServerCommon) resetBulkTransportAndWait(ctx context.Context, transport *TransportConn, req BulkResetRequest) error {
	if s == nil {
		return errBulkServerNil
	}
	err := runBulkRecoveryTaskContext(ctx, newServerBulkResetRecoveryTask(s, transport.logicalConnSnapshot(), transport, req))
	if err != nil && bulkRecoveryErrorRetryable(err) {
		s.bestEffortBulkResetTransport(transport, req)
	}
	return err
}

func (s *ServerCommon) cleanupBulkTransportReset(ctx context.Context, transport *TransportConn, req BulkResetRequest, wait bool) error {
	if wait {
		return s.resetBulkTransportAndWait(ctx, transport, req)
	}
	s.bestEffortBulkResetTransport(transport, req)
	return nil
}

func (c *ClientCommon) handleBulkRecoveryOverflow(epoch uint64, req BulkResetRequest) {
	route := c.clientSessionRouteSnapshot()
	route.epoch = epoch
	c.handleBulkRecoveryOverflowAtRoute(route, req)
}

func (c *ClientCommon) handleBulkRecoveryOverflowAtRoute(route clientSessionRoute, req BulkResetRequest) {
	if c == nil {
		return
	}
	err := fmt.Errorf("%w: bulk=%s data=%d", errBulkRecoveryQueueFull, req.BulkID, req.DataID)
	c.reportBulkRecoveryError(err)
	if c.clientSessionRouteCurrent(route) && route.epoch != 0 {
		c.stopClientSessionIfCurrent(route.epoch, "bulk recovery queue full", err)
	}
}

func (s *ServerCommon) handleBulkRecoveryOverflow(logical *LogicalConn, transport *TransportConn, req BulkResetRequest) {
	if s == nil {
		return
	}
	err := fmt.Errorf("%w: bulk=%s data=%d", errBulkRecoveryQueueFull, req.BulkID, req.DataID)
	s.reportBulkRecoveryError(err)
	if logical != nil && transport != nil && transport.IsCurrent() {
		s.detachLogicalSessionTransport(logical, "bulk recovery queue full", err)
	}
}

func (c *ClientCommon) reportBulkRecoveryError(err error) {
	if c == nil || err == nil {
		return
	}
	c.mu.Lock()
	debug := c.showError || c.debugMode
	c.mu.Unlock()
	if debug {
		fmt.Printf("notify bulk reset recovery failed: %v\n", err)
	}
}

func (s *ServerCommon) reportBulkRecoveryError(err error) {
	if s == nil || err == nil {
		return
	}
	s.mu.RLock()
	debug := s.showError || s.debugMode
	s.mu.RUnlock()
	if debug {
		fmt.Printf("notify bulk reset recovery failed: %v\n", err)
	}
}