subscriber.go

// Copyright 2025 samber.
//
// 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
//
// https://github.com/samber/ro/blob/main/licenses/LICENSE.apache.md
//
// 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 ro

import (
	"context"
	"sync"
	"sync/atomic"
)

// Subscriber implements the Observer and Subscription interfaces. While the Observer is
// the public API for consuming the values of an Observable, all Observers get
// converted to a Subscriber, in order to provide Subscription-like capabilities
// such as `Unsubscribe()`. Subscriber is a common type in samber/ro, and crucial for
// implementing operators, but it is rarely used as a public API.
type Subscriber[T any] interface {
	Subscription
	Observer[T]
}

var _ Subscriber[int] = (*subscriberImpl[int])(nil)

// NewSubscriber creates a new Subscriber from an Observer. If the Observer
// is already a Subscriber, it is returned as is. Otherwise, a new Subscriber
// is created that wraps the Observer.
//
// The returned Subscriber will unsubscribe from the destination Observer when
// Unsubscribe() is called.
//
// This method is safe for concurrent use.
//
// It is rarely used as a public API.
func NewSubscriber[T any](destination Observer[T]) Subscriber[T] {
	return NewSafeSubscriber(destination)
}

// NewSafeSubscriber creates a new Subscriber from an Observer. If the Observer
// is already a Subscriber, it is returned as is. Otherwise, a new Subscriber
// is created that wraps the Observer.
//
// The returned Subscriber will unsubscribe from the destination Observer when
// Unsubscribe() is called.
//
// This method is safe for concurrent use.
//
// It is rarely used as a public API.
func NewSafeSubscriber[T any](destination Observer[T]) Subscriber[T] {
	return NewSubscriberWithConcurrencyMode(destination, ConcurrencyModeSafe)
}

// NewUnsafeSubscriber creates a new Subscriber from an Observer. If the Observer
// is already a Subscriber, it is returned as is. Otherwise, a new Subscriber
// is created that wraps the Observer.
//
// The returned Subscriber will unsubscribe from the destination Observer when
// Unsubscribe() is called.
//
// This method is not safe for concurrent use.
//
// It is rarely used as a public API.
func NewUnsafeSubscriber[T any](destination Observer[T]) Subscriber[T] {
	return NewSubscriberWithConcurrencyMode(destination, ConcurrencyModeUnsafe)
}

// NewEventuallySafeSubscriber creates a new Subscriber from an Observer. If the Observer
// is already a Subscriber, it is returned as is. Otherwise, a new Subscriber
// is created that wraps the Observer.
//
// The returned Subscriber will unsubscribe from the destination Observer when
// Unsubscribe() is called.
//
// This method is safe for concurrent use, but concurrent messages are dropped.
//
// It is rarely used as a public API.
func NewEventuallySafeSubscriber[T any](destination Observer[T]) Subscriber[T] {
	return NewSubscriberWithConcurrencyMode(destination, ConcurrencyModeEventuallySafe)
}

// NewSubscriberWithConcurrencyMode creates a new Subscriber from an Observer. If the Observer
// is already a Subscriber, it is returned as is. Otherwise, a new Subscriber
// is created that wraps the Observer.
//
// The returned Subscriber will unsubscribe from the destination Observer when
// Unsubscribe() is called.
//
// It is rarely used as a public API.
func NewSubscriberWithConcurrencyMode[T any](destination Observer[T], mode ConcurrencyMode) Subscriber[T] {
	// Spinlock is ignored because it is too slow when chaining operators. Spinlock should be used
	// only for short-lived local locks.
	switch mode {
	case ConcurrencyModeSafe:
		return newSubscriberImpl(mode, false, BackpressureBlock, destination)
	case ConcurrencyModeUnsafe:
		return newSubscriberImpl(mode, true, BackpressureBlock, destination)
	case ConcurrencyModeEventuallySafe:
		return newSubscriberImpl(mode, false, BackpressureDrop, destination)
	default:
		panic("invalid concurrency mode")
	}
}

// newSubscriberImpl creates a new subscriber implementation with the specified
// synchronization behavior and destination observer.
func newSubscriberImpl[T any](mode ConcurrencyMode, noLock bool, backpressure Backpressure, destination Observer[T]) Subscriber[T] {
	// Protect against multiple encapsulation layers.
	if subscriber, ok := destination.(Subscriber[T]); ok {
		return subscriber
	}

	subscriber := &subscriberImpl[T]{
		status:       0, // KindNext
		backpressure: backpressure,

		noLock:      noLock,
		destination: destination,

		Subscription: NewSubscription(nil),
		mode:         mode,
	}

	if subscription, ok := destination.(Subscription); ok {
		subscription.Add(subscriber.Unsubscribe)
	}

	return subscriber
}

type subscriberImpl[T any] struct {
	// While mutex is used for synchronization of producer, status is used for storing state of
	// the subscriber. Using the mutex for reading the status would have create a dead lock if
	// an Observer calls Unsubscribe(), IsClosed(), HasThrown(), IsCompleted() synchronously.
	//
	// 0 - KindNext
	// 1 - KindError
	// 2 - KindComplete
	status       int32
	backpressure Backpressure

	_ [59]byte // padding to prevent false sharing

	// Mutexes are much faster than channels.
	//
	// A concrete sync.Mutex (instead of an xsync.Mutex interface) keeps the
	// lock fast path inlinable: interface dispatch on every Next is measurably
	// slower. ConcurrencyModeUnsafe skips the lock via the noLock flag.
	//
	// Also, generators has been added in go1.23. A different implem of Observable/Observer
	// might reduce latency induced by mutexes.
	//
	// It could be interesting to implement a lock-free version of this,
	// with message drop instead of backpressure, and when SLO must be kept under
	// control (real-time streams?).
	mu          sync.Mutex
	noLock      bool
	destination Observer[T]

	Subscription

	mode ConcurrencyMode
}

func (s *subscriberImpl[T]) lock() {
	if !s.noLock {
		s.mu.Lock()
	}
}

func (s *subscriberImpl[T]) unlock() {
	if !s.noLock {
		s.mu.Unlock()
	}
}

func (s *subscriberImpl[T]) tryLock() bool {
	if s.noLock {
		return true
	}

	return s.mu.TryLock()
}

// Implements Observer.
func (s *subscriberImpl[T]) Next(v T) {
	s.NextWithContext(context.Background(), v)
}

// Implements Observer.
func (s *subscriberImpl[T]) NextWithContext(ctx context.Context, v T) {
	if s.destination == nil {
		return
	}

	if s.backpressure == BackpressureDrop {
		if !s.tryLock() {
			OnDroppedNotification(ctx, NewNotificationNext(v))
			return
		}
	} else {
		s.lock()
	}

	if atomic.LoadInt32(&s.status) == 0 {
		s.destination.NextWithContext(ctx, v)
	} else {
		OnDroppedNotification(ctx, NewNotificationNext(v))
	}

	s.unlock()
}

// Implements Observer.
func (s *subscriberImpl[T]) Error(err error) {
	s.ErrorWithContext(context.Background(), err)
}

// Implements Observer.
func (s *subscriberImpl[T]) ErrorWithContext(ctx context.Context, err error) {
	s.lock()

	if atomic.CompareAndSwapInt32(&s.status, 0, 1) {
		if s.destination != nil {
			s.destination.ErrorWithContext(ctx, err)
		}
	} else {
		OnDroppedNotification(ctx, NewNotificationError[T](err))
	}

	s.unlock()

	s.unsubscribe()
}

// Implements Observer.
func (s *subscriberImpl[T]) Complete() {
	s.CompleteWithContext(context.Background())
}

// Implements Observer.
func (s *subscriberImpl[T]) CompleteWithContext(ctx context.Context) {
	s.lock()

	if atomic.CompareAndSwapInt32(&s.status, 0, 2) {
		if s.destination != nil {
			s.destination.CompleteWithContext(ctx)
		}
	} else {
		OnDroppedNotification(ctx, NewNotificationComplete[T]())
	}

	s.unlock()

	s.unsubscribe()
}

// Implements Observer.
func (s *subscriberImpl[T]) IsClosed() bool {
	return atomic.LoadInt32(&s.status) != 0
}

// Implements Observer.
func (s *subscriberImpl[T]) HasThrown() bool {
	return atomic.LoadInt32(&s.status) == 1
}

// Implements Observer.
func (s *subscriberImpl[T]) IsCompleted() bool {
	return atomic.LoadInt32(&s.status) == 2
}

// Implements Observer.
func (s *subscriberImpl[T]) Unsubscribe() {
	if atomic.CompareAndSwapInt32(&s.status, 0, 2) {
		s.unsubscribe()
	}
}

func (s *subscriberImpl[T]) unsubscribe() {
	// s.Subscription.Unsubscribe() is protected against concurrent calls.
	s.Subscription.Unsubscribe()
}