notify/control_batch_sender.go

package notify

import (
	"bytes"
	"context"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

const (
	controlBatchMaxPayloads           = 16
	controlBatchMaxPayloadBytes       = 32 * 1024 * 1024
	controlBatchMaxQueuedBytes        = 32 * 1024 * 1024
	controlBatchCriticalReservedBytes = 1 * 1024 * 1024
	controlBatchMaxCriticalBurst      = 16
)

type controlPriority uint8

const (
	controlPriorityNormal controlPriority = iota
	controlPriorityCritical
)

const (
	controlBatchRequestQueued int32 = iota
	controlBatchRequestStarted
	controlBatchRequestCanceled
)

type controlBatchRequestState struct {
	value atomic.Int32
}

type controlBatchRequest struct {
	ctx       context.Context
	payload   []byte
	writeTime time.Duration
	priority  controlPriority
	done      chan error
	state     *controlBatchRequestState
	queueSize int64
}

type controlBatchSender struct {
	binding    *transportBinding
	normalCh   chan controlBatchRequest
	criticalCh chan controlBatchRequest
	stopCh     chan struct{}
	doneCh     chan struct{}
	budgetWake chan struct{}
	stopCtx    context.Context
	stopCancel context.CancelFunc

	stopOnce        sync.Once
	flushMu         sync.Mutex
	admissionMu     sync.Mutex
	admitting       sync.WaitGroup
	admissionClosed bool
	queued          atomic.Int64
	queuedSize      atomic.Int64
	errMu           sync.Mutex
	err             error
}

func newControlBatchSender(binding *transportBinding) *controlBatchSender {
	stopCtx, stopCancel := context.WithCancel(context.Background())
	sender := &controlBatchSender{
		binding:    binding,
		normalCh:   make(chan controlBatchRequest, controlBatchMaxPayloads*4),
		criticalCh: make(chan controlBatchRequest, controlBatchMaxPayloads*2),
		stopCh:     make(chan struct{}),
		doneCh:     make(chan struct{}),
		budgetWake: make(chan struct{}, 1),
		stopCtx:    stopCtx,
		stopCancel: stopCancel,
	}
	go sender.run()
	return sender
}

func (s *controlBatchSender) submit(payload []byte, writeTimeout time.Duration) error {
	return s.submitContext(context.Background(), payload, writeTimeout, controlPriorityNormal)
}

func (s *controlBatchSender) submitContext(ctx context.Context, payload []byte, writeTimeout time.Duration, priority controlPriority) error {
	if s == nil {
		return newTransportSendError(TransportSendStageTransport, errTransportDetached)
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := s.errSnapshot(); err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		return newTransportSendError(TransportSendStageQueue, err)
	}
	if len(payload) > controlBatchMaxPayloadBytes {
		return newTransportSendError(
			TransportSendStageQueue,
			fmt.Errorf("control payload is %d bytes, maximum is %d; use stream or bulk transfer", len(payload), controlBatchMaxPayloadBytes),
		)
	}
	req := controlBatchRequest{
		ctx:       ctx,
		payload:   payload,
		writeTime: maxDuration(0, writeTimeout),
		priority:  priority,
		queueSize: int64(maxInt(1, len(payload))),
	}
	if submitted, err := s.tryDirectSubmit(req); submitted {
		return err
	}
	queueCtx, cancelQueue := contextWithTimeoutUpperBound(ctx, req.writeTime)
	defer cancelQueue()
	req.ctx = queueCtx
	if req.ctx.Done() != nil {
		req.payload = bytes.Clone(payload)
	}

	s.queued.Add(1)
	if err := s.reserveQueueBytes(req.ctx, req.queueSize, priority); err != nil {
		s.queued.Add(-1)
		return err
	}
	req.done = make(chan error, 1)
	req.state = &controlBatchRequestState{}
	queued := false
	defer func() {
		if !queued {
			s.queued.Add(-1)
			s.releaseQueueBytes(req.queueSize)
		}
	}()
	if !s.beginAdmission() {
		return s.stoppedErr()
	}
	select {
	case <-req.ctx.Done():
		s.endAdmission()
		return newTransportSendError(TransportSendStageQueue, req.ctx.Err())
	case <-s.stopCh:
		s.endAdmission()
		return s.stoppedErr()
	case s.requestChannel(priority) <- req:
		queued = true
		s.endAdmission()
	}
	select {
	case err := <-req.done:
		return err
	case <-s.stopCh:
		if req.tryCancel() {
			return s.stoppedErr()
		}
		if req.ctx != nil && req.ctx.Done() != nil {
			return s.stoppedErr()
		}
		return <-req.done
	case <-req.ctx.Done():
		if req.tryCancel() {
			return newTransportSendError(TransportSendStageQueue, req.ctx.Err())
		}
		return newTransportSendError(TransportSendStageQueue, req.ctx.Err())
	}
}

func contextWithTimeoutUpperBound(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
	if ctx == nil {
		ctx = context.Background()
	}
	if timeout <= 0 {
		return ctx, func() {}
	}
	deadline := time.Now().Add(timeout)
	if current, ok := ctx.Deadline(); ok && !deadline.Before(current) {
		return ctx, func() {}
	}
	return context.WithDeadline(ctx, deadline)
}

func (s *controlBatchSender) tryDirectSubmit(req controlBatchRequest) (bool, error) {
	if s == nil {
		return true, newTransportSendError(TransportSendStageTransport, errTransportDetached)
	}
	if err := s.errSnapshot(); err != nil {
		return true, err
	}
	select {
	case <-req.ctx.Done():
		return true, newTransportSendError(TransportSendStageQueue, req.ctx.Err())
	case <-s.stopCh:
		return true, s.stoppedErr()
	default:
	}
	if req.ctx.Done() != nil {
		return false, nil
	}
	if s.queued.Load() != 0 || !s.flushMu.TryLock() {
		return false, nil
	}
	defer s.flushMu.Unlock()
	if s.queued.Load() != 0 {
		return false, nil
	}
	if err := s.errSnapshot(); err != nil {
		return true, err
	}
	if err := s.flushDirect(req); err != nil {
		if stage, ok := TransportSendErrorStage(err); ok && stage == TransportSendStageQueue {
			return true, err
		}
		s.markFailed(err)
		s.waitAdmissions()
		s.failPending(err, nil, nil)
		return true, err
	}
	return true, nil
}

func (s *controlBatchSender) run() {
	defer close(s.doneCh)
	var pendingNormal []controlBatchRequest
	var pendingCritical []controlBatchRequest
	criticalBurst := 0
	for {
		forceNormal := criticalBurst >= controlBatchMaxCriticalBurst
		req, ok := s.nextRequest(&pendingNormal, &pendingCritical, forceNormal)
		if !ok {
			s.waitAdmissions()
			s.failPending(s.stoppedErr(), pendingNormal, pendingCritical)
			return
		}
		normalWaiting := s.hasNormalRequest(&pendingNormal)
		if req.priority == controlPriorityCritical {
			switch {
			case normalWaiting && criticalBurst >= controlBatchMaxCriticalBurst:
				pushControlPendingFront(req, &pendingNormal, &pendingCritical)
				var available bool
				req, available = s.tryNextRequest(controlPriorityNormal, &pendingNormal, &pendingCritical)
				if !available {
					s.waitAdmissions()
					s.failPending(s.stoppedErr(), pendingNormal, pendingCritical)
					return
				}
				forceNormal = true
			case !normalWaiting:
				criticalBurst = 0
			}
		}
		batch := []controlBatchRequest{req}
		batchBytes := len(req.payload)
		softLimit := s.batchSoftPayloadLimit()
		batchLimit := controlBatchMaxPayloads
		if req.priority == controlPriorityCritical && normalWaiting {
			remaining := controlBatchMaxCriticalBurst - criticalBurst
			if remaining < batchLimit {
				batchLimit = remaining
			}
		}
		for len(batch) < batchLimit {
			next, available := s.tryNextRequest(req.priority, &pendingNormal, &pendingCritical)
			if !available {
				break
			}
			if !controlBatchCanAppend(batchBytes, len(next.payload), softLimit) {
				pushControlPendingFront(next, &pendingNormal, &pendingCritical)
				break
			}
			batch = append(batch, next)
			batchBytes += len(next.payload)
		}

		s.flushMu.Lock()
		if req.priority == controlPriorityNormal && !forceNormal && criticalBurst < controlBatchMaxCriticalBurst {
			if critical, available := s.tryNextRequest(controlPriorityCritical, &pendingNormal, &pendingCritical); available {
				pendingNormal = append(append(make([]controlBatchRequest, 0, len(batch)+len(pendingNormal)), batch...), pendingNormal...)
				batch = []controlBatchRequest{critical}
				batchBytes = len(critical.payload)
				batchLimit = controlBatchMaxCriticalBurst - criticalBurst
				if batchLimit > controlBatchMaxPayloads {
					batchLimit = controlBatchMaxPayloads
				}
				for len(batch) < batchLimit {
					next, ok := s.tryNextRequest(controlPriorityCritical, &pendingNormal, &pendingCritical)
					if !ok {
						break
					}
					if !controlBatchCanAppend(batchBytes, len(next.payload), softLimit) {
						pushControlPendingFront(next, &pendingNormal, &pendingCritical)
						break
					}
					batch = append(batch, next)
					batchBytes += len(next.payload)
				}
			}
		}
		batchPriority := batch[0].priority
		err := s.flushQueued(batch)
		s.flushMu.Unlock()
		if err != nil {
			s.markFailed(err)
			s.waitAdmissions()
			s.failPending(err, pendingNormal, pendingCritical)
			return
		}
		if batchPriority == controlPriorityCritical {
			criticalBurst += len(batch)
			if criticalBurst > controlBatchMaxCriticalBurst {
				criticalBurst = controlBatchMaxCriticalBurst
			}
		} else {
			criticalBurst = 0
		}
	}
}

func (s *controlBatchSender) nextRequest(pendingNormal *[]controlBatchRequest, pendingCritical *[]controlBatchRequest, forceNormal bool) (controlBatchRequest, bool) {
	if forceNormal {
		if req, ok := popControlPending(pendingNormal); ok {
			return req, true
		}
		select {
		case req := <-s.normalCh:
			return req, true
		default:
		}
	}
	if req, ok := popControlPending(pendingCritical); ok {
		return req, true
	}
	select {
	case req := <-s.criticalCh:
		return req, true
	default:
	}
	if req, ok := popControlPending(pendingNormal); ok {
		return req, true
	}
	select {
	case <-s.stopCh:
		return controlBatchRequest{}, false
	case req := <-s.criticalCh:
		return req, true
	case req := <-s.normalCh:
		return req, true
	}
}

func (s *controlBatchSender) tryNextRequest(priority controlPriority, pendingNormal *[]controlBatchRequest, pendingCritical *[]controlBatchRequest) (controlBatchRequest, bool) {
	pending := pendingNormal
	if priority == controlPriorityCritical {
		pending = pendingCritical
	}
	if req, ok := popControlPending(pending); ok {
		return req, true
	}
	select {
	case <-s.stopCh:
		return controlBatchRequest{}, false
	case req := <-s.requestChannel(priority):
		return req, true
	default:
		return controlBatchRequest{}, false
	}
}

func popControlPending(pending *[]controlBatchRequest) (controlBatchRequest, bool) {
	if pending == nil || len(*pending) == 0 {
		return controlBatchRequest{}, false
	}
	req := (*pending)[0]
	*pending = (*pending)[1:]
	return req, true
}

func pushControlPendingFront(req controlBatchRequest, pendingNormal *[]controlBatchRequest, pendingCritical *[]controlBatchRequest) {
	pending := pendingNormal
	if req.priority == controlPriorityCritical {
		pending = pendingCritical
	}
	*pending = append([]controlBatchRequest{req}, (*pending)...)
}

func (s *controlBatchSender) requestChannel(priority controlPriority) chan controlBatchRequest {
	if priority == controlPriorityCritical {
		return s.criticalCh
	}
	return s.normalCh
}

func (s *controlBatchSender) hasNormalRequest(pendingNormal *[]controlBatchRequest) bool {
	if pendingNormal != nil && len(*pendingNormal) > 0 {
		return true
	}
	return s != nil && len(s.normalCh) > 0
}

func controlBatchCanAppend(batchBytes int, nextBytes int, softLimit int) bool {
	if batchBytes == 0 {
		return true
	}
	if softLimit <= 0 {
		return false
	}
	return nextBytes <= softLimit-batchBytes
}

func (s *controlBatchSender) batchSoftPayloadLimit() int {
	if s == nil || s.binding == nil {
		return controlAdaptiveSoftPayloadFallbackBytes
	}
	return s.binding.controlAdaptiveSoftPayloadBytesSnapshot()
}

func (s *controlBatchSender) reserveQueueBytes(ctx context.Context, size int64, priority controlPriority) error {
	limit := int64(controlBatchMaxQueuedBytes)
	if priority == controlPriorityCritical {
		limit += int64(controlBatchCriticalReservedBytes)
	}
	for {
		current := s.queuedSize.Load()
		if (current == 0 && size > limit) || size <= limit-current {
			if s.queuedSize.CompareAndSwap(current, current+size) {
				return nil
			}
			continue
		}
		select {
		case <-ctx.Done():
			return newTransportSendError(TransportSendStageQueue, ctx.Err())
		case <-s.stopCh:
			return s.stoppedErr()
		case <-s.budgetWake:
		}
	}
}

func (s *controlBatchSender) releaseQueueBytes(size int64) {
	if s == nil || size <= 0 {
		return
	}
	s.queuedSize.Add(-size)
	select {
	case s.budgetWake <- struct{}{}:
	default:
	}
}

func (s *controlBatchSender) flushDirect(req controlBatchRequest) error {
	if s == nil || s.binding == nil {
		return newTransportSendError(TransportSendStageTransport, errTransportDetached)
	}
	queue := s.binding.queueSnapshot()
	if queue == nil {
		return newTransportSendError(TransportSendStageTransport, errTransportFrameQueueUnavailable)
	}
	var preWriteErr error
	didWrite := false
	started := time.Now()
	lockAcquired, err := s.binding.withConnWriteLockContextStopDeadlineManaged(req.ctx, s.stopCh, writeDeadlineFromTimeout(req.writeTime), func(conn net.Conn) error {
		if stoppedErr := s.errSnapshot(); stoppedErr != nil {
			preWriteErr = stoppedErr
			return nil
		}
		if ctxErr := req.contextErr(); ctxErr != nil {
			preWriteErr = ctxErr
			return nil
		}
		didWrite = true
		return writeFramedPayloadBatchUnlocked(conn, queue, [][]byte{req.payload})
	})
	if preWriteErr != nil {
		return preWriteErr
	}
	if !lockAcquired {
		if ctxErr := req.contextErr(); ctxErr != nil {
			return ctxErr
		}
		if stoppedErr := s.errSnapshot(); stoppedErr != nil {
			return stoppedErr
		}
		return newTransportSendError(TransportSendStageTransport, err)
	}
	if didWrite {
		s.binding.observeControlAdaptivePayloadWrite(len(req.payload), time.Since(started), req.writeTime, err)
	}
	if err != nil {
		if lockAcquired {
			// A framed write may have left a partial frame on the wire. Do not
			// let a later request reuse this connection.
			s.binding.closeConn()
		}
		return newTransportSendError(TransportSendStageWrite, err)
	}
	return nil
}

func (s *controlBatchSender) flushQueued(requests []controlBatchRequest) error {
	if s == nil || s.binding == nil {
		err := newTransportSendError(TransportSendStageTransport, errTransportDetached)
		s.finishRequests(requests, err)
		return err
	}
	queue := s.binding.queueSnapshot()
	if queue == nil {
		err := newTransportSendError(TransportSendStageTransport, errTransportFrameQueueUnavailable)
		s.finishRequests(requests, err)
		return err
	}
	pending := requests
	for len(pending) > 0 {
		pending = s.finishCanceledRequests(pending)
		if len(pending) == 0 {
			return nil
		}
		if stoppedErr := s.errSnapshot(); stoppedErr != nil {
			s.finishRequests(pending, stoppedErr)
			return stoppedErr
		}

		waitCtx, cleanupWait := s.controlBatchWaitContext(pending)
		active := make([]controlBatchRequest, 0, len(pending))
		var preWriteErr error
		payloadBytes := 0
		didWrite := false
		started := time.Now()
		writeDeadline := earlierWriteDeadline(
			writeDeadlineFromTimeout(controlBatchRequestsShortestWriteTimeout(pending)),
			controlBatchRequestsEarliestDeadline(pending),
		)
		lockAcquired, err := s.binding.withConnWriteLockContextStopDeadlineManaged(waitCtx, s.stopCh, writeDeadline, func(conn net.Conn) error {
			if stoppedErr := s.errSnapshot(); stoppedErr != nil {
				preWriteErr = stoppedErr
				return nil
			}
			for _, item := range pending {
				if cancelErr, canceled := item.cancelBeforeStart(); canceled {
					s.finishRequest(item, cancelErr)
					continue
				}
				if !item.tryStart() {
					s.finishRequest(item, item.canceledErr())
					continue
				}
				active = append(active, item)
				payloadBytes += len(item.payload)
			}
			if len(active) == 0 {
				return nil
			}
			payloads := make([][]byte, 0, len(active))
			for _, item := range active {
				payloads = append(payloads, item.payload)
			}
			didWrite = true
			return writeFramedPayloadBatchUnlocked(conn, queue, payloads)
		})
		cleanupWait()
		if preWriteErr != nil {
			s.finishRequests(pending, preWriteErr)
			return preWriteErr
		}
		if !lockAcquired {
			remaining := s.finishCanceledRequests(pending)
			if len(remaining) < len(pending) {
				pending = remaining
				continue
			}
			if stoppedErr := s.errSnapshot(); stoppedErr != nil {
				s.finishRequests(pending, stoppedErr)
				return stoppedErr
			}
			transportErr := newTransportSendError(TransportSendStageTransport, err)
			s.finishRequests(pending, transportErr)
			return transportErr
		}
		if len(active) == 0 {
			if err == nil {
				return nil
			}
			writeErr := newTransportSendError(TransportSendStageWrite, err)
			s.finishRequests(pending, writeErr)
			return writeErr
		}
		writeTimeout := controlBatchRequestsShortestWriteTimeout(active)
		if didWrite {
			s.binding.observeControlAdaptivePayloadWrite(payloadBytes, time.Since(started), writeTimeout, err)
		}
		if err != nil {
			if lockAcquired {
				// A framed write may have left a partial frame on the wire. Do not
				// let a later request reuse this connection.
				s.binding.closeConn()
			}
			err = newTransportSendError(TransportSendStageWrite, err)
		}
		s.finishRequests(active, err)
		return err
	}
	return nil
}

func (s *controlBatchSender) controlBatchWaitContext(requests []controlBatchRequest) (context.Context, func()) {
	base := context.Background()
	if s != nil && s.stopCtx != nil {
		base = s.stopCtx
	}
	ctx, cancel := context.WithCancel(base)
	stops := make([]func() bool, 0, len(requests))
	for _, item := range requests {
		if item.ctx == nil || item.ctx.Done() == nil {
			continue
		}
		stops = append(stops, context.AfterFunc(item.ctx, cancel))
	}
	return ctx, func() {
		for _, stop := range stops {
			stop()
		}
		cancel()
	}
}

func controlBatchRequestsShortestWriteTimeout(batch []controlBatchRequest) time.Duration {
	var timeout time.Duration
	for _, item := range batch {
		if item.writeTime <= 0 {
			continue
		}
		if timeout == 0 || item.writeTime < timeout {
			timeout = item.writeTime
		}
	}
	return timeout
}

func controlBatchRequestsEarliestDeadline(batch []controlBatchRequest) time.Time {
	var deadline time.Time
	for _, item := range batch {
		candidate := contextDeadline(item.ctx)
		deadline = earlierWriteDeadline(deadline, candidate)
	}
	return deadline
}

func (s *controlBatchSender) finishCanceledRequests(requests []controlBatchRequest) []controlBatchRequest {
	remaining := requests[:0]
	for _, item := range requests {
		if cancelErr, canceled := item.cancelBeforeStart(); canceled {
			s.finishRequest(item, cancelErr)
			continue
		}
		remaining = append(remaining, item)
	}
	return remaining
}

func (s *controlBatchSender) finishRequests(requests []controlBatchRequest, err error) {
	for _, item := range requests {
		s.finishRequest(item, err)
	}
}

func (s *controlBatchSender) finishRequest(req controlBatchRequest, err error) {
	if s != nil {
		s.queued.Add(-1)
		s.releaseQueueBytes(req.queueSize)
	}
	req.done <- err
}

func (s *controlBatchSender) beginAdmission() bool {
	if s == nil {
		return false
	}
	s.admissionMu.Lock()
	defer s.admissionMu.Unlock()
	if s.admissionClosed {
		return false
	}
	s.admitting.Add(1)
	return true
}

func (s *controlBatchSender) endAdmission() {
	if s != nil {
		s.admitting.Done()
	}
}

func (s *controlBatchSender) waitAdmissions() {
	if s != nil {
		s.admitting.Wait()
	}
}

func (s *controlBatchSender) stop() {
	if s == nil {
		return
	}
	s.markFailed(newTransportSendError(TransportSendStageTransport, errTransportDetached))
	<-s.doneCh
	s.flushMu.Lock()
	s.flushMu.Unlock()
}

func (s *controlBatchSender) failPending(err error, pendingNormal []controlBatchRequest, pendingCritical []controlBatchRequest) {
	for _, item := range pendingCritical {
		s.finishRequest(item, err)
	}
	for _, item := range pendingNormal {
		s.finishRequest(item, err)
	}
	for {
		select {
		case item := <-s.criticalCh:
			s.finishRequest(item, err)
		case item := <-s.normalCh:
			s.finishRequest(item, err)
		default:
			return
		}
	}
}

func (s *controlBatchSender) setErr(err error) {
	if s == nil || err == nil {
		return
	}
	s.errMu.Lock()
	if s.err == nil {
		s.err = err
	}
	s.errMu.Unlock()
}

func (s *controlBatchSender) markFailed(err error) {
	if s == nil {
		return
	}
	s.setErr(err)
	s.stopOnce.Do(func() {
		s.admissionMu.Lock()
		s.admissionClosed = true
		if s.stopCancel != nil {
			s.stopCancel()
		}
		close(s.stopCh)
		s.admissionMu.Unlock()
	})
}

func (s *controlBatchSender) errSnapshot() error {
	if s == nil {
		return newTransportSendError(TransportSendStageTransport, errTransportDetached)
	}
	s.errMu.Lock()
	defer s.errMu.Unlock()
	return s.err
}

func (s *controlBatchSender) stoppedErr() error {
	if err := s.errSnapshot(); err != nil {
		return err
	}
	return newTransportSendError(TransportSendStageTransport, errTransportDetached)
}

func (r controlBatchRequest) tryStart() bool {
	if r.state == nil {
		return true
	}
	return r.state.value.CompareAndSwap(controlBatchRequestQueued, controlBatchRequestStarted)
}

func (r controlBatchRequest) tryCancel() bool {
	return r.state != nil && r.state.value.CompareAndSwap(controlBatchRequestQueued, controlBatchRequestCanceled)
}

func (r controlBatchRequest) cancelBeforeStart() (error, bool) {
	if r.state == nil {
		if err := r.contextErr(); err != nil {
			return err, true
		}
		return nil, false
	}
	if r.state.value.Load() == controlBatchRequestCanceled {
		return r.canceledErr(), true
	}
	if err := r.contextErr(); err != nil && r.tryCancel() {
		return err, true
	}
	if r.state.value.Load() == controlBatchRequestCanceled {
		return r.canceledErr(), true
	}
	return nil, false
}

func (r controlBatchRequest) contextErr() error {
	if r.ctx == nil {
		return nil
	}
	err := r.ctx.Err()
	if err == nil {
		if deadline, ok := r.ctx.Deadline(); ok && !time.Now().Before(deadline) {
			err = context.DeadlineExceeded
		}
	}
	return newTransportSendError(TransportSendStageQueue, err)
}

func (r controlBatchRequest) canceledErr() error {
	if r.ctx != nil && r.ctx.Err() != nil {
		return newTransportSendError(TransportSendStageQueue, r.ctx.Err())
	}
	return newTransportSendError(TransportSendStageQueue, context.Canceled)
}

func maxInt(left int, right int) int {
	if left > right {
		return left
	}
	return right
}

func maxDuration(left time.Duration, right time.Duration) time.Duration {
	if left > right {
		return left
	}
	return right
}