l-ai-chat-client-go/subscription.go

package chatclient

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

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	pb "gitea.fe.kaien.ru/kaien/l-ai-chat-proto"
)

type Delivery struct {
	Message *pb.Message
	Replay  bool

	subscriptionID string
	generation     uint64
	sub            *Subscription
	token          *deliveryToken
}

type Handler func(context.Context, Delivery) error

type Subscription struct {
	client *Client
	topic  string

	ctx    context.Context
	cancel context.CancelFunc

	deliveries chan Delivery
	done       chan struct{}
	ackNotify  chan struct{}

	mu             sync.Mutex
	err            error
	closed         bool
	generation     uint64
	subscriptionID string
	tracker        ackTracker
	ackTarget      int64
	ackConfirmed   int64

	readyStream pb.ChatService_SubscribeClient
}

func newSubscription(ctx context.Context, c *Client, req SubscribeRequest) (*Subscription, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	subCtx, cancel := context.WithCancel(ctx)
	s := &Subscription{
		client:     c,
		topic:      req.TopicID,
		ctx:        subCtx,
		cancel:     cancel,
		deliveries: make(chan Delivery, c.opts.subscriptionBufferSize),
		done:       make(chan struct{}),
		ackNotify:  make(chan struct{}, 1),
		tracker:    newACKTracker(),
	}

	stream, err := s.openReady(subCtx)
	if err != nil {
		cancel()
		return nil, err
	}
	s.readyStream = stream
	s.generation = 1

	go s.run()
	go s.runACK()
	return s, nil
}

func (s *Subscription) Deliveries() <-chan Delivery { return s.deliveries }
func (s *Subscription) Done() <-chan struct{}       { return s.done }

func (s *Subscription) Err() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.err
}

func (s *Subscription) Close() error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil
	}
	s.closed = true
	s.mu.Unlock()
	s.cancel()

	timer := time.NewTimer(s.client.opts.closeTimeout)
	defer timer.Stop()
	select {
	case <-s.done:
		return nil
	case <-timer.C:
		return context.DeadlineExceeded
	}
}

func (s *Subscription) MarkProcessed(ctx context.Context, d Delivery) error {
	if ctx == nil {
		ctx = context.Background()
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}
	if d.sub != s || d.token == nil {
		return ErrUnknownDelivery
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosed
	}
	if d.generation != s.generation {
		return ErrStaleDelivery
	}
	boundary, ok, _ := s.tracker.mark(d.token)
	if !ok {
		return ErrUnknownDelivery
	}
	if boundary > s.ackTarget {
		s.ackTarget = boundary
		s.signalACKLocked()
	}
	return nil
}

func (s *Subscription) Run(ctx context.Context, handler Handler) error {
	if ctx == nil {
		ctx = context.Background()
	}
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case d, ok := <-s.deliveries:
			if !ok {
				return ignoreClosed(s.Err())
			}
			if err := handler(ctx, d); err != nil {
				return err
			}
			if err := s.MarkProcessed(ctx, d); err != nil {
				return err
			}
		}
	}
}

func (s *Subscription) run() {
	defer func() {
		s.cancel()
		close(s.deliveries)
		close(s.done)
		s.client.unregister(s)
	}()

	stream := s.readyStream
	s.readyStream = nil
	backoff := s.client.opts.reconnectInitial

	for {
		err := s.recvUntilError(stream)
		if s.ctx.Err() != nil {
			s.setErr(nil)
			return
		}
		if isReplacement(err) {
			s.setErr(fmt.Errorf("%w: %v", ErrSubscriptionReplaced, err))
			return
		}
		if !isRetryableStreamError(err) {
			s.setErr(err)
			return
		}
		if err := s.client.opts.sleeper(s.ctx, s.client.opts.jitter(backoff)); err != nil {
			s.setErr(nil)
			return
		}
		backoff = nextBackoff(backoff, s.client.opts.reconnectMax)

		next, err := s.openReady(s.ctx)
		if err != nil {
			if s.ctx.Err() != nil {
				s.setErr(nil)
				return
			}
			if isReplacement(err) {
				s.setErr(fmt.Errorf("%w: %v", ErrSubscriptionReplaced, err))
				return
			}
			if !isRetryableStreamError(err) {
				s.setErr(err)
				return
			}
			continue
		}
		s.resetGeneration()
		stream = next
		backoff = s.client.opts.reconnectInitial
	}
}

func (s *Subscription) recvUntilError(stream pb.ChatService_SubscribeClient) error {
	for {
		dm, err := stream.Recv()
		if err != nil {
			return err
		}
		if dm.GetReady() {
			return fmt.Errorf("%w: unexpected ready marker after subscription registration", ErrProtocolViolation)
		}
		if dm.GetMessage() == nil {
			return fmt.Errorf("%w: delivery message is missing message", ErrProtocolViolation)
		}
		token := &deliveryToken{}
		delivery := Delivery{
			Message:        dm.GetMessage(),
			Replay:         dm.GetReplay(),
			subscriptionID: dm.GetSubscriptionId(),
			token:          token,
			sub:            s,
		}
		s.mu.Lock()
		if s.closed {
			s.mu.Unlock()
			return ErrClosed
		}
		if s.subscriptionID == "" {
			s.subscriptionID = dm.GetSubscriptionId()
		}
		if dm.GetSubscriptionId() == "" || dm.GetSubscriptionId() != s.subscriptionID {
			s.mu.Unlock()
			return fmt.Errorf("%w: inconsistent subscription_id", ErrProtocolViolation)
		}
		delivery.generation = s.generation
		s.tracker.add(dm.GetMessage().GetSequence(), token)
		s.mu.Unlock()

		select {
		case s.deliveries <- delivery:
		case <-s.ctx.Done():
			return ErrClosed
		}
	}
}

func (s *Subscription) openReady(ctx context.Context) (pb.ChatService_SubscribeClient, error) {
	stream, err := s.client.api.Subscribe(ctx, &pb.SubscribeRequest{TopicId: s.topic})
	if err != nil {
		return nil, err
	}
	dm, err := stream.Recv()
	if err != nil {
		return nil, err
	}
	if !dm.GetReady() || dm.GetMessage() != nil {
		return nil, fmt.Errorf("%w: subscribe stream did not start with ready marker", ErrProtocolViolation)
	}
	return stream, nil
}

func (s *Subscription) resetGeneration() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.generation++
	s.subscriptionID = ""
	s.tracker = newACKTracker()
	s.ackTarget = 0
	s.ackConfirmed = 0
}

func (s *Subscription) runACK() {
	for {
		select {
		case <-s.ctx.Done():
			return
		case <-s.ackNotify:
		}

		for {
			gen, subID, target, ok := s.nextACK()
			if !ok {
				break
			}
			resp, err := s.client.api.Acknowledge(s.ctx, &pb.AcknowledgeRequest{
				TopicId:         s.topic,
				LastAckSequence: target,
				SubscriptionId:  subID,
			})
			if err != nil {
				current := s.sameGeneration(gen)
				if s.ctx.Err() != nil || !isRetryableACKError(err) || !current {
					if s.ctx.Err() == nil && current && !isRetryableACKError(err) {
						s.setErr(err)
						s.cancel()
					}
					break
				}
				_ = s.client.opts.sleeper(s.ctx, s.client.opts.jitter(s.client.opts.reconnectInitial))
				continue
			}
			s.confirmACK(gen, resp.GetLastAckSequence())
		}
	}
}

func (s *Subscription) nextACK() (uint64, string, int64, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed || s.subscriptionID == "" || s.ackTarget <= s.ackConfirmed {
		return 0, "", 0, false
	}
	return s.generation, s.subscriptionID, s.ackTarget, true
}

func (s *Subscription) sameGeneration(gen uint64) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.generation == gen && !s.closed
}

func (s *Subscription) confirmACK(gen uint64, confirmed int64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.generation != gen || s.closed {
		return
	}
	if confirmed > s.ackConfirmed {
		s.ackConfirmed = confirmed
		s.tracker.removeConfirmed(confirmed)
	}
	if s.ackTarget > s.ackConfirmed {
		s.signalACKLocked()
	}
}

func (s *Subscription) signalACKLocked() {
	select {
	case s.ackNotify <- struct{}{}:
	default:
	}
}

func (s *Subscription) setErr(err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.err == nil {
		s.err = err
	}
	s.closed = true
}

func nextBackoff(cur, max time.Duration) time.Duration {
	if cur <= 0 {
		return max
	}
	next := cur * 2
	if next < cur || next > max {
		return max
	}
	return next
}

func isReplacement(err error) bool {
	return status.Code(err) == codes.Aborted
}

func isRetryableStreamError(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, io.EOF) {
		return true
	}
	switch status.Code(err) {
	case codes.Unavailable, codes.ResourceExhausted, codes.DeadlineExceeded:
		return true
	case codes.Canceled:
		return false
	default:
		return false
	}
}

func isRetryableACKError(err error) bool {
	switch status.Code(err) {
	case codes.Unavailable, codes.ResourceExhausted, codes.DeadlineExceeded:
		return true
	default:
		return false
	}
}