iris/context_wrapper.go

package iris

import (
	"sync"
	"time"

	"github.com/kataras/iris/v12/context"
)

type (
	// ContextSetter is an interface which can be implemented by a struct
	// to set the iris.Context to the struct.
	// The receiver must be a pointer of the struct.
	ContextSetter interface {
		// SetContext sets the iris.Context to the struct.
		SetContext(Context)
	}

	// ContextSetterPtr is a pointer of T which implements the `ContextSetter` interface.
	// The T must be a struct.
	ContextSetterPtr[T any] interface {
		*T
		ContextSetter
	}

	// emptyContextSetter is an empty struct which implements the `ContextSetter` interface.
	emptyContextSetter struct{}
)

// SetContext method implements `ContextSetter` interface.
func (*emptyContextSetter) SetContext(Context) {}

// ContextPool is a pool of T. It's used to acquire and release custom context.
// Use of custom implementation or `NewContextPool`.
//
// See `NewContextWrapper` and `NewContextPool` for more.
type (
	ContextPool[T any] interface {
		// Acquire must return a new T from a pool.
		Acquire(ctx Context) T
		// Release must put the T back to the pool.
		Release(T)
	}

	// syncContextPool is a sync pool implementation of T.
	// It's used to acquire and release T.
	// The contextPtr is acquired from the sync pool and released back to the sync pool after the handler's execution.
	// The contextPtr is passed to the handler as an argument.
	// ThecontextPtr is not shared between requests.
	// The contextPtr must implement the `ContextSetter` interface.
	// The T must be a struct.
	// The contextPtr must be a pointer of T.
	syncContextPool[T any, contextPtr ContextSetterPtr[T]] struct {
		pool *sync.Pool
	}
)

// Ensure that syncContextPool implements ContextPool.
var _ ContextPool[*emptyContextSetter] = (*syncContextPool[emptyContextSetter, *emptyContextSetter])(nil)

// NewContextPool returns a new ContextPool default implementation which
// uses sync.Pool to implement its Acquire and Release methods.
// The contextPtr is acquired from the sync pool and released back to the sync pool after the handler's execution.
// The contextPtr is passed to the handler as an argument.
// ThecontextPtr is not shared between requests.
// The contextPtr must implement the `ContextSetter` interface.
// The T must be a struct.
// The contextPtr must be a pointer of T.
//
// Example:
// w := iris.NewContextWrapper(iris.NewContextPool[myCustomContext, *myCustomContext]())
func NewContextPool[T any, contextPtr ContextSetterPtr[T]]() ContextPool[contextPtr] {
	return &syncContextPool[T, contextPtr]{
		pool: &sync.Pool{
			New: func() interface{} {
				var t contextPtr = new(T)
				return t
			},
		},
	}
}

// Acquire returns a new T from the sync pool.
func (p *syncContextPool[T, contextPtr]) Acquire(ctx Context) contextPtr {
	// var t contextPtr
	// if v := p.pool.Get(); v == nil {
	// 	t = new(T)
	// } else {
	// 	t = v.(contextPtr)
	// }

	t := p.pool.Get().(contextPtr)
	t.SetContext(ctx)
	return t
}

// Release puts the T back to the sync pool.
func (p *syncContextPool[T, contextPtr]) Release(t contextPtr) {
	p.pool.Put(t)
}

// ContextWrapper is a wrapper for handlers which expect a T instead of iris.Context.
//
// See the `NewContextWrapper` function for more.
type ContextWrapper[T any] struct {
	pool ContextPool[T]
}

// NewContextWrapper returns a new ContextWrapper.
// If pool is nil, a default pool is used.
// The default pool's AcquireFunc returns a zero value of T.
// The default pool's ReleaseFunc does nothing.
// The default pool is used when the pool is nil.
// Use the `iris.NewContextPool[T, *T]()` to pass a simple context pool.
// Then, use the `Handler` method to wrap custom handlers to iris ones.
//
// Example: https://github.com/kataras/iris/tree/main/_examples/routing/custom-context
func NewContextWrapper[T any](pool ContextPool[T]) *ContextWrapper[T] {
	if pool == nil {
		panic("pool cannot be nil")
	}

	return &ContextWrapper[T]{
		pool: pool,
	}
}

// Pool returns the pool, useful when manually Acquire and Release of custom context is required.
func (w *ContextWrapper[T]) Pool() ContextPool[T] {
	return w.pool
}

// Handler wraps the handler with the pool's Acquire and Release methods.
// It returns a new handler which expects a T instead of iris.Context.
// The T is the type of the pool.
// The T is acquired from the pool and released back to the pool after the handler's execution.
// The T is passed to the handler as an argument.
// The T is not shared between requests.
func (w *ContextWrapper[T]) Handler(handler func(T)) Handler {
	if handler == nil {
		return nil
	}

	return func(ctx Context) {
		newT := w.pool.Acquire(ctx)
		handler(newT)
		w.pool.Release(newT)
	}
}

// Handlers wraps the handlers with the pool's Acquire and Release methods.
func (w *ContextWrapper[T]) Handlers(handlers ...func(T)) context.Handlers {
	newHandlers := make(context.Handlers, len(handlers))
	for i, handler := range handlers {
		newHandlers[i] = w.Handler(handler)
	}

	return newHandlers
}

// HandlerReturnError same as `Handler` but it converts a handler which returns an error.
func (w *ContextWrapper[T]) HandlerReturnError(handler func(T) error) func(Context) error {
	if handler == nil {
		return nil
	}

	return func(ctx Context) error {
		newT := w.pool.Acquire(ctx)
		err := handler(newT)
		w.pool.Release(newT)
		return err
	}
}

// HandlerReturnDuration same as `Handler` but it converts a handler which returns a time.Duration.
func (w *ContextWrapper[T]) HandlerReturnDuration(handler func(T) time.Duration) func(Context) time.Duration {
	if handler == nil {
		return nil
	}

	return func(ctx Context) time.Duration {
		newT := w.pool.Acquire(ctx)
		duration := handler(newT)
		w.pool.Release(newT)
		return duration
	}
}

// Filter same as `Handler` but it converts a handler to Filter.
func (w *ContextWrapper[T]) Filter(handler func(T) bool) Filter {
	if handler == nil {
		return nil
	}

	return func(ctx Context) bool {
		newT := w.pool.Acquire(ctx)
		shouldContinue := handler(newT)
		w.pool.Release(newT)
		return shouldContinue
	}
}

// FallbackViewFunc same as `Handler` but it converts a handler to FallbackViewFunc.
func (w *ContextWrapper[T]) FallbackViewFunc(handler func(ctx T, err ErrViewNotExist) error) FallbackViewFunc {
	if handler == nil {
		return nil
	}

	return func(ctx Context, err ErrViewNotExist) error {
		newT := w.pool.Acquire(ctx)
		returningErr := handler(newT, err)
		w.pool.Release(newT)
		return returningErr
	}
}