notify/control_batch_sender.go

package notify

import (
	"net"
	"sync"
	"time"
)

const controlBatchMaxPayloads = 16

type controlBatchRequest struct {
	payload  []byte
	deadline time.Time
	done     chan error
}

type controlBatchSender struct {
	binding *transportBinding
	reqCh   chan controlBatchRequest
	stopCh  chan struct{}
	doneCh  chan struct{}

	stopOnce sync.Once
	errMu    sync.Mutex
	err      error
}

func newControlBatchSender(binding *transportBinding) *controlBatchSender {
	sender := &controlBatchSender{
		binding: binding,
		reqCh:   make(chan controlBatchRequest, controlBatchMaxPayloads*4),
		stopCh:  make(chan struct{}),
		doneCh:  make(chan struct{}),
	}
	go sender.run()
	return sender
}

func (s *controlBatchSender) submit(payload []byte, deadline time.Time) error {
	if s == nil {
		return errTransportDetached
	}
	req := controlBatchRequest{
		payload:  payload,
		deadline: deadline,
		done:     make(chan error, 1),
	}
	if err := s.errSnapshot(); err != nil {
		return err
	}
	select {
	case <-s.stopCh:
		return s.stoppedErr()
	case s.reqCh <- req:
	}
	return <-req.done
}

func (s *controlBatchSender) run() {
	defer close(s.doneCh)
	for {
		req, ok := s.nextRequest()
		if !ok {
			return
		}
		batch := []controlBatchRequest{req}
	drain:
		for len(batch) < controlBatchMaxPayloads {
			select {
			case <-s.stopCh:
				s.failPending(s.stoppedErr())
				return
			case next := <-s.reqCh:
				batch = append(batch, next)
			default:
				break drain
			}
		}
		payloads := make([][]byte, 0, len(batch))
		for _, item := range batch {
			payloads = append(payloads, item.payload)
		}
		err := s.flush(payloads, controlBatchRequestsEarliestDeadline(batch))
		if err != nil {
			s.setErr(err)
			for _, item := range batch {
				item.done <- err
			}
			s.failPending(err)
			return
		}
		for _, item := range batch {
			item.done <- nil
		}
	}
}

func (s *controlBatchSender) nextRequest() (controlBatchRequest, bool) {
	select {
	case <-s.stopCh:
		s.failPending(s.stoppedErr())
		return controlBatchRequest{}, false
	case req := <-s.reqCh:
		return req, true
	}
}

func controlBatchRequestsEarliestDeadline(batch []controlBatchRequest) time.Time {
	var deadline time.Time
	for _, item := range batch {
		if item.deadline.IsZero() {
			continue
		}
		if deadline.IsZero() || item.deadline.Before(deadline) {
			deadline = item.deadline
		}
	}
	return deadline
}

func (s *controlBatchSender) flush(payloads [][]byte, deadline time.Time) error {
	if s == nil || s.binding == nil {
		return errTransportDetached
	}
	queue := s.binding.queueSnapshot()
	if queue == nil {
		return errTransportFrameQueueUnavailable
	}
	return s.binding.withConnWriteLockDeadline(deadline, func(conn net.Conn) error {
		return writeFramedPayloadBatchUnlocked(conn, queue, payloads)
	})
}

func (s *controlBatchSender) stop() {
	if s == nil {
		return
	}
	s.stopOnce.Do(func() {
		s.setErr(errTransportDetached)
		close(s.stopCh)
	})
	<-s.doneCh
}

func (s *controlBatchSender) failPending(err error) {
	for {
		select {
		case item := <-s.reqCh:
			item.done <- 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) errSnapshot() error {
	if s == nil {
		return 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 errTransportDetached
}