go-libs/pkg/messaging/publish/circuit/circuit_breaker.go

package circuit

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"time"

	"github.com/ThreeDotsLabs/watermill/message"
	"github.com/google/uuid"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"

	"github.com/hanzo-fi/go-libs/v5/pkg/messaging/publish/circuit/storage"
	logging "github.com/hanzo-fi/go-libs/v5/pkg/observe/log"
)

type State string

const (
	// StateOpen is the state when the circuit breaker is open and does not allow
	// requests to pass through. Instead, it writes the message in the database
	// and waits for the "openInterval" to pass before switching to the
	// "half-open" state.
	StateOpen State = "open"
	// StateClosed is the default state. It allows requests to pass through.
	StateClose State = "close"
)

type CircuitBreaker struct {
	logger logging.Logger

	publisher message.Publisher
	store     storage.Store

	stateMu sync.RWMutex
	state   State

	stopOnce sync.Once
	doneOnce sync.Once

	publisherCloseOnce sync.Once
	publisherCloseErr  error

	loopMu      sync.Mutex
	loopCancel  context.CancelFunc
	loopStarted bool

	// openInterval is the time interval for the "open" state, before switching
	// to the "half-open" state.
	openInterval      time.Duration
	openIntervalTimer *time.Timer

	sendChan    chan *internalMessage
	stopChannel chan struct{}
	stopped     chan struct{}
}

type internalMessage struct {
	topic string
	msg   *message.Message

	errChan chan error
}

func NewCircuitBreaker(
	logger logging.Logger,
	publisher message.Publisher,
	store storage.Store,
	openIntervalDuration time.Duration,
) *CircuitBreaker {
	return &CircuitBreaker{
		stopChannel: make(chan struct{}),
		stopped:     make(chan struct{}),
		logger:      logger,
		publisher:   publisher,
		store:       store,

		state: StateClose,

		openInterval:      openIntervalDuration,
		openIntervalTimer: time.NewTimer(openIntervalDuration),

		// no capacity, we want to block the loop if the sendChan is not consumed
		sendChan: make(chan *internalMessage, 1),
	}
}

func (cb *CircuitBreaker) GetState() State {
	cb.stateMu.RLock()
	defer cb.stateMu.RUnlock()
	return cb.state
}

func (cb *CircuitBreaker) setState(state State) {
	cb.stateMu.Lock()
	cb.state = state
	cb.stateMu.Unlock()
}

func (cb *CircuitBreaker) OpenState() {
	cb.setState(StateOpen)
	cb.openIntervalTimer.Reset(cb.openInterval)

	cb.logger.Info("Circuit breaker switched to the open state")
}

func (cb *CircuitBreaker) HalfOpenState() {
	cb.openIntervalTimer.Stop()
	cb.logger.Info("Circuit breaker switched to the half open state")
}

func (cb *CircuitBreaker) CloseState() {
	cb.setState(StateClose)
	cb.openIntervalTimer.Stop()

	cb.logger.Info("Circuit breaker switched to the close state")
}

func (cb *CircuitBreaker) Loop(ctx context.Context) {
	loopCtx, cancel := context.WithCancel(ctx)
	if !cb.startLoop(cancel) {
		cancel()
		return
	}
	defer func() {
		cb.clearLoopCancel()
		cancel()
		cb.markStopped()
	}()

	select {
	case <-loopCtx.Done():
		return
	case <-cb.stopChannel:
		return
	default:
	}

	// Start in the half open state to fetch the messages from the database
	cb.HalfOpenState()
	if err := cb.catchUpDatabase(loopCtx, ctx); err != nil {
		if loopCtx.Err() != nil {
			return
		}
		cb.OpenState()
		// Don't switch to closed state if there was an error
	} else {
		// Only switch to closed state if catchup was successful
		cb.CloseState()
	}

	for {
		select {
		case <-loopCtx.Done():
			return

		case <-cb.stopChannel:
			return

		case <-cb.openIntervalTimer.C:
			// openInterval passed, let's switch to the half-open state

			cb.HalfOpenState()

			if err := cb.catchUpDatabase(loopCtx, ctx); err != nil {
				if loopCtx.Err() != nil {
					return
				}
				cb.OpenState()
				continue
			}

			// we successfully published the messages, let's switch to the closed state
			cb.CloseState()

		case msg, ok := <-cb.sendChan:
			if !ok {
				// sendChan closed
				return
			}

			switch cb.GetState() {
			case StateClose:
				// We are in the closed state, send the message to the publisher

				cb.logger.Info("Circular breaker is in the closed state, sending the message to the publisher")
				err := cb.publisher.Publish(msg.topic, msg.msg)
				if err != nil {
					// error publishing the message, let's switch to the open state
					cb.OpenState()

					cb.logger.Info("Failed to publish the message, switching to the open state")
					// write the message in the database
					err = cb.store.Insert(ctx, msg.topic, msg.msg.Payload, msg.msg.Metadata)
					if err != nil {
						msg.errChan <- err
						continue
					}
				}
			case StateOpen:
				// We are in the open state, write the message in the database
				cb.logger.Info("Circuit breaker is in the open state, writing the message in the database")

				err := cb.store.Insert(ctx, msg.topic, msg.msg.Payload, msg.msg.Metadata)
				if err != nil {
					msg.errChan <- err
					continue
				}
			}

			msg.errChan <- nil
		}
	}
}

func (cb *CircuitBreaker) catchUpDatabase(loopCtx context.Context, storeCtx context.Context) error {
	for {
		if err := loopCtx.Err(); err != nil {
			return err
		}

		// fetch the oldest messages from the database
		messages, err := cb.store.List(loopCtx)
		if err != nil {
			// error fetching messages, let's switch back to the open state
			return err
		}

		if len(messages) == 0 {
			return nil
		}

		messagesToDelete := make([]uint64, 0)
		var publishError error
		var stopError error
		for _, msg := range messages {
			if err := loopCtx.Err(); err != nil {
				stopError = err
				break
			}

			if err := storeCtx.Err(); err != nil {
				return err
			}

			// We need to publish the messages one by one in order to know
			// which one failed.

			message, err := newMessage(storeCtx, msg.Data, msg.Metadata)
			if err != nil {
				publishError = err
				break
			}

			err = cb.publisher.Publish(msg.Topic, message)
			if err != nil {
				publishError = err
				break
			}

			messagesToDelete = append(messagesToDelete, msg.ID)
		}

		if len(messagesToDelete) > 0 {
			err = cb.store.Delete(storeCtx, messagesToDelete)
			if err != nil {
				// error deleting messages, let's switch back to the open state
				return err
			}
		}

		if publishError != nil {
			// we failed to publish all the messages
			return publishError
		}
		if stopError != nil {
			return stopError
		}
	}
}

func (cb *CircuitBreaker) Publish(topic string, messages ...*message.Message) error {
	for _, msg := range messages {
		select {
		case <-cb.stopChannel:
			return errors.New("circuit breaker closed")
		default:
		}

		errChan := make(chan error, 1)

		internalMessage := &internalMessage{
			topic:   topic,
			msg:     msg,
			errChan: errChan,
		}

		select {
		case cb.sendChan <- internalMessage:
		case <-cb.stopChannel:
			return errors.New("circuit breaker closed")
		case <-cb.stopped:
			return errors.New("circuit breaker closed")
		}

		if err := waitAcceptedMessageResult(errChan, cb.stopped); err != nil {
			return err
		}
	}

	return nil
}

func waitAcceptedMessageResult(errChan <-chan error, stopped <-chan struct{}) error {
	select {
	case err := <-errChan:
		return err
	case <-stopped:
		select {
		case err := <-errChan:
			return err
		default:
			return errors.New("circuit breaker closed")
		}
	}
}

func (cb *CircuitBreaker) Close() error {
	cb.stopOnce.Do(func() {
		close(cb.stopChannel)
		cb.cancelLoop()
	})

	if cb.hasLoopStarted() {
		<-cb.stopped
	} else {
		cb.markStopped()
	}

	cb.publisherCloseOnce.Do(func() {
		cb.publisherCloseErr = cb.publisher.Close()
	})

	return cb.publisherCloseErr
}

func (cb *CircuitBreaker) startLoop(cancel context.CancelFunc) bool {
	cb.loopMu.Lock()
	defer cb.loopMu.Unlock()

	if cb.loopStarted {
		return false
	}

	cb.loopStarted = true
	cb.loopCancel = cancel
	return true
}

func (cb *CircuitBreaker) clearLoopCancel() {
	cb.loopMu.Lock()
	cb.loopCancel = nil
	cb.loopMu.Unlock()
}

func (cb *CircuitBreaker) cancelLoop() {
	cb.loopMu.Lock()
	cancel := cb.loopCancel
	cb.loopMu.Unlock()

	if cancel != nil {
		cancel()
	}
}

func (cb *CircuitBreaker) hasLoopStarted() bool {
	cb.loopMu.Lock()
	defer cb.loopMu.Unlock()
	return cb.loopStarted
}

func (cb *CircuitBreaker) markStopped() {
	cb.doneOnce.Do(func() {
		close(cb.stopped)
	})
}

const (
	otelContextKey = "otel-context"
)

func newMessage(ctx context.Context, data []byte, metadata map[string]string) (*message.Message, error) {
	if metadata == nil {
		metadata = make(map[string]string)
	}

	msg := message.NewMessage(uuid.NewString(), data)

	otelContext, ok := metadata[otelContextKey]
	if ok {
		carrier := propagation.MapCarrier{}
		err := json.Unmarshal([]byte(otelContext), &carrier)
		if err != nil {
			return nil, err
		}
		ctx = otel.GetTextMapPropagator().Extract(ctx, carrier)
	}

	msg.SetContext(ctx)
	for k, v := range metadata {
		msg.Metadata.Set(k, v)
	}

	return msg, nil
}