psrpc/internal/bus/bus_kv.go

// Copyright 2023 LiveKit, Inc.
//
// 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 bus

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"math/rand"
	"sync"
	"time"

	"github.com/gammazero/deque"
	"github.com/hanzokv/go/v9"
	"github.com/zeebo/xxh3"
	"go.uber.org/multierr"
	"golang.org/x/exp/maps"
	"golang.org/x/exp/slices"
	"google.golang.org/protobuf/proto"

	"github.com/hanzoai/psrpc/internal/logger"
)

const (
	lockExpiration          = time.Second * 5
	reconcilerRetryInterval = time.Second
	minReadRetryInterval    = time.Millisecond * 100
	maxReadRetryInterval    = time.Second
	publishBuckets          = 17
)

type kvMessageBus struct {
	rc  kv.UniversalClient
	ctx context.Context
	ps  *kv.PubSub

	mu     sync.Mutex
	subs   map[string]*kvSubList
	queues map[string]*kvSubList

	wakeup          chan struct{}
	ops             *kvWriteOpQueue
	dirtyChannels   map[string]struct{}
	currentChannels map[string]struct{}

	publishQueues [publishBuckets]*kvPublishQueue
}

func NewKVMessageBus(rc kv.UniversalClient) MessageBus {
	ctx := context.Background()
	r := &kvMessageBus{
		rc:     rc,
		ctx:    ctx,
		ps:     rc.Subscribe(ctx),
		subs:   map[string]*kvSubList{},
		queues: map[string]*kvSubList{},

		wakeup:          make(chan struct{}, 1),
		ops:             &kvWriteOpQueue{},
		dirtyChannels:   map[string]struct{}{},
		currentChannels: map[string]struct{}{},
	}
	for i := range len(r.publishQueues) {
		r.publishQueues[i] = newKVPublishQueue(r.ctx, r.rc)
	}
	go r.readWorker()
	go r.writeWorker()
	return r
}

func (r *kvMessageBus) Publish(_ context.Context, channel Channel, msg proto.Message) error {
	b, err := serialize(msg, "")
	if err != nil {
		return err
	}

	bucket := xxh3.HashString(channel.Legacy) % publishBuckets
	r.publishQueues[bucket].Enqueue(channel.Legacy, b)
	return nil
}

func (r *kvMessageBus) Subscribe(ctx context.Context, channel Channel, size int) (Reader, error) {
	return r.subscribe(ctx, channel.Legacy, size, r.subs, false)
}

func (r *kvMessageBus) SubscribeQueue(ctx context.Context, channel Channel, size int) (Reader, error) {
	return r.subscribe(ctx, channel.Legacy, size, r.queues, true)
}

func (r *kvMessageBus) subscribe(ctx context.Context, channel string, size int, subLists map[string]*kvSubList, queue bool) (Reader, error) {
	ctx, cancel := context.WithCancel(ctx)
	sub := &kvSubscription{
		bus:     r,
		ctx:     ctx,
		cancel:  cancel,
		channel: channel,
		msgChan: make(chan *kv.Message, size),
		queue:   queue,
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	subList, ok := subLists[channel]
	if !ok {
		subList = &kvSubList{}
		subLists[channel] = subList
		r.reconcileSubscriptions(channel)
	}
	subList.subs = append(subList.subs, sub)

	return sub, nil
}

func (r *kvMessageBus) unsubscribe(channel string, queue bool, sub *kvSubscription) {
	r.mu.Lock()
	defer r.mu.Unlock()

	var subLists map[string]*kvSubList
	if queue {
		subLists = r.queues
	} else {
		subLists = r.subs
	}

	subList, ok := subLists[channel]
	if !ok {
		return
	}
	i := slices.Index(subList.subs, sub)
	if i == -1 {
		return
	}

	subList.subs = slices.Delete(subList.subs, i, i+1)

	if len(subList.subs) == 0 {
		delete(subLists, channel)
		r.reconcileSubscriptions(channel)
	}
}

func (r *kvMessageBus) readWorker() {
	var delay time.Duration
	for {
		msg, err := r.ps.ReceiveMessage(r.ctx)
		if err != nil {
			logger.Error(err, "kv receive message failed")

			time.Sleep(delay)
			if delay *= 2; delay == 0 {
				delay = minReadRetryInterval
			} else if delay > maxReadRetryInterval {
				delay = maxReadRetryInterval
			}
			continue
		}
		delay = 0

		r.mu.Lock()
		if subList, ok := r.subs[msg.Channel]; ok {
			subList.dispatch(msg)
		}
		if subList, ok := r.queues[msg.Channel]; ok {
			subList.dispatchQueue(msg)
		}
		r.mu.Unlock()
	}
}

func (r *kvMessageBus) reconcileSubscriptions(channel string) {
	r.dirtyChannels[channel] = struct{}{}
	r.enqueueWriteOp(&kvReconcileSubscriptionsOp{r})
}

func (r *kvMessageBus) enqueueWriteOp(op kvWriteOp) {
	r.ops.push(op)
	select {
	case r.wakeup <- struct{}{}:
	default:
	}
}

func (r *kvMessageBus) writeWorker() {
	for range r.wakeup {
		r.ops.drain()
	}
}

// ----------------------------------------------

type kvWriteOpQueue struct {
	mu  sync.Mutex
	ops deque.Deque[kvWriteOp]
}

func (q *kvWriteOpQueue) empty() bool {
	q.mu.Lock()
	defer q.mu.Unlock()
	return q.ops.Len() == 0
}

func (q *kvWriteOpQueue) push(op kvWriteOp) {
	q.mu.Lock()
	defer q.mu.Unlock()
	q.ops.PushBack(op)
}

func (q *kvWriteOpQueue) drain() {
	q.mu.Lock()
	for q.ops.Len() > 0 {
		op := q.ops.PopFront()
		q.mu.Unlock()
		if err := op.run(); err != nil {
			logger.Error(err, "kv write message failed")
		}
		q.mu.Lock()
	}
	q.mu.Unlock()
}

//-----------------------------------------------------

type kvWriteOp interface {
	run() error
}

// ----------------------------------------------------

type kvReconcileSubscriptionsOp struct {
	*kvMessageBus
}

func (r *kvReconcileSubscriptionsOp) run() error {
	r.mu.Lock()
	for len(r.dirtyChannels) > 0 {
		subscribe := make(map[string]struct{}, len(r.dirtyChannels))
		unsubscribe := make(map[string]struct{}, len(r.dirtyChannels))
		for c := range r.dirtyChannels {
			_, current := r.currentChannels[c]
			desired := r.subs[c] != nil || r.queues[c] != nil
			if !current && desired {
				subscribe[c] = struct{}{}
			} else if current && !desired {
				unsubscribe[c] = struct{}{}
			}
		}
		maps.Clear(r.dirtyChannels)
		r.mu.Unlock()

		var subscribeErr, unsubscribeErr error
		if len(subscribe) != 0 {
			subscribeErr = r.ps.Subscribe(r.ctx, maps.Keys(subscribe)...)
		}
		if len(unsubscribe) != 0 {
			unsubscribeErr = r.ps.Unsubscribe(r.ctx, maps.Keys(unsubscribe)...)
		}

		if err := multierr.Combine(subscribeErr, unsubscribeErr); err != nil {
			logger.Error(err, "kv subscription reconciliation failed")
			time.Sleep(reconcilerRetryInterval)
		}

		r.mu.Lock()
		if subscribeErr != nil {
			maps.Copy(r.dirtyChannels, subscribe)
		} else {
			maps.Copy(r.currentChannels, subscribe)
		}
		if unsubscribeErr != nil {
			maps.Copy(r.dirtyChannels, unsubscribe)
		} else {
			for c := range unsubscribe {
				delete(r.currentChannels, c)
			}
		}
	}
	r.mu.Unlock()
	return nil
}

// ----------------------------------------------------

type kvSubList struct {
	subs []*kvSubscription
	next int
}

func (r *kvSubList) dispatchQueue(msg *kv.Message) {
	if r.next >= len(r.subs) {
		r.next = 0
	}
	r.subs[r.next].write(msg)
	r.next++
}

func (r *kvSubList) dispatch(msg *kv.Message) {
	for _, sub := range r.subs {
		sub.write(msg)
	}
}

// ----------------------------------------------------

type kvSubscription struct {
	bus     *kvMessageBus
	ctx     context.Context
	cancel  context.CancelFunc
	channel string
	msgChan chan *kv.Message
	queue   bool
}

func (r *kvSubscription) write(msg *kv.Message) {
	select {
	case r.msgChan <- msg:
	case <-r.ctx.Done():
	}
}

func (r *kvSubscription) read() ([]byte, bool) {
	for {
		var msg *kv.Message
		var ok bool
		select {
		case msg, ok = <-r.msgChan:
			if !ok {
				return nil, false
			}
		case <-r.ctx.Done():
			return nil, false
		}

		if r.queue {
			sha := sha256.Sum256([]byte(msg.Payload))
			hash := base64.StdEncoding.EncodeToString(sha[:])
			acquired, err := r.bus.rc.SetNX(r.ctx, hash, rand.Int(), lockExpiration).Result()
			if err != nil || !acquired {
				continue
			}
		}

		return []byte(msg.Payload), true
	}
}

func (r *kvSubscription) Close() error {
	r.cancel()
	r.bus.unsubscribe(r.channel, r.queue, r)
	close(r.msgChan)
	return nil
}

// ----------------------------------------------------

type kvPublishMessage struct {
	channel string
	payload []byte
}

type kvPublishQueue struct {
	ctx context.Context
	rc  kv.UniversalClient

	lock     sync.Mutex
	messages []kvPublishMessage
	wakeup   chan struct{}
}

func newKVPublishQueue(ctx context.Context, rc kv.UniversalClient) *kvPublishQueue {
	r := &kvPublishQueue{
		ctx:    ctx,
		rc:     rc,
		wakeup: make(chan struct{}, 1),
	}

	go r.worker()
	return r
}

func (r *kvPublishQueue) Enqueue(channel string, payload []byte) {
	r.lock.Lock()
	defer r.lock.Unlock()

	r.messages = append(r.messages, kvPublishMessage{channel, payload})
	select {
	case r.wakeup <- struct{}{}:
	default:
	}
}

func (r *kvPublishQueue) worker() {
	for {
		select {
		case <-r.wakeup:
		case <-r.ctx.Done():
			return
		}

		r.lock.Lock()
		messages := r.messages
		r.messages = nil
		r.lock.Unlock()

		// using a pipeline to handle kv servers with a high RTT
		// (https://redis.io/docs/latest/develop/using-commands/pipelining/).
		//
		// This is doing oppotunistic batching + pipelining.
		// When a message is published, this worker is signalled immediately and
		// the message will be sent immediately. While the pipeline is executing,
		// messages will get queued as the pipeline execution takes one round trip
		// exchange with the server. The next round will batch all those queued
		// messages. For small RTTs, this will send messages without any delay.
		pipeline := r.rc.Pipeline()
		for _, msg := range messages {
			pipeline.Publish(r.ctx, msg.channel, msg.payload)
		}
		cmds, err := pipeline.Exec(r.ctx)
		if err != nil {
			logger.Error(err, "pipeline execution failed", "numCommands", len(cmds))
		}
	}
}

// ----------------------------------------------------