context_wrapper.go 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215
  1. package iris
  2. import (
  3. "sync"
  4. "time"
  5. "github.com/kataras/iris/v12/context"
  6. )
  7. type (
  8. // ContextSetter is an interface which can be implemented by a struct
  9. // to set the iris.Context to the struct.
  10. // The receiver must be a pointer of the struct.
  11. ContextSetter interface {
  12. // SetContext sets the iris.Context to the struct.
  13. SetContext(Context)
  14. }
  15. // ContextSetterPtr is a pointer of T which implements the `ContextSetter` interface.
  16. // The T must be a struct.
  17. ContextSetterPtr[T any] interface {
  18. *T
  19. ContextSetter
  20. }
  21. // emptyContextSetter is an empty struct which implements the `ContextSetter` interface.
  22. emptyContextSetter struct{}
  23. )
  24. // SetContext method implements `ContextSetter` interface.
  25. func (*emptyContextSetter) SetContext(Context) {}
  26. // ContextPool is a pool of T. It's used to acquire and release custom context.
  27. // Use of custom implementation or `NewContextPool`.
  28. //
  29. // See `NewContextWrapper` and `NewContextPool` for more.
  30. type (
  31. ContextPool[T any] interface {
  32. // Acquire must return a new T from a pool.
  33. Acquire(ctx Context) T
  34. // Release must put the T back to the pool.
  35. Release(T)
  36. }
  37. // syncContextPool is a sync pool implementation of T.
  38. // It's used to acquire and release T.
  39. // The contextPtr is acquired from the sync pool and released back to the sync pool after the handler's execution.
  40. // The contextPtr is passed to the handler as an argument.
  41. // ThecontextPtr is not shared between requests.
  42. // The contextPtr must implement the `ContextSetter` interface.
  43. // The T must be a struct.
  44. // The contextPtr must be a pointer of T.
  45. syncContextPool[T any, contextPtr ContextSetterPtr[T]] struct {
  46. pool *sync.Pool
  47. }
  48. )
  49. // Ensure that syncContextPool implements ContextPool.
  50. var _ ContextPool[*emptyContextSetter] = (*syncContextPool[emptyContextSetter, *emptyContextSetter])(nil)
  51. // NewContextPool returns a new ContextPool default implementation which
  52. // uses sync.Pool to implement its Acquire and Release methods.
  53. // The contextPtr is acquired from the sync pool and released back to the sync pool after the handler's execution.
  54. // The contextPtr is passed to the handler as an argument.
  55. // ThecontextPtr is not shared between requests.
  56. // The contextPtr must implement the `ContextSetter` interface.
  57. // The T must be a struct.
  58. // The contextPtr must be a pointer of T.
  59. //
  60. // Example:
  61. // w := iris.NewContextWrapper(iris.NewContextPool[myCustomContext, *myCustomContext]())
  62. func NewContextPool[T any, contextPtr ContextSetterPtr[T]]() ContextPool[contextPtr] {
  63. return &syncContextPool[T, contextPtr]{
  64. pool: &sync.Pool{
  65. New: func() interface{} {
  66. var t contextPtr = new(T)
  67. return t
  68. },
  69. },
  70. }
  71. }
  72. // Acquire returns a new T from the sync pool.
  73. func (p *syncContextPool[T, contextPtr]) Acquire(ctx Context) contextPtr {
  74. // var t contextPtr
  75. // if v := p.pool.Get(); v == nil {
  76. // t = new(T)
  77. // } else {
  78. // t = v.(contextPtr)
  79. // }
  80. t := p.pool.Get().(contextPtr)
  81. t.SetContext(ctx)
  82. return t
  83. }
  84. // Release puts the T back to the sync pool.
  85. func (p *syncContextPool[T, contextPtr]) Release(t contextPtr) {
  86. p.pool.Put(t)
  87. }
  88. // ContextWrapper is a wrapper for handlers which expect a T instead of iris.Context.
  89. //
  90. // See the `NewContextWrapper` function for more.
  91. type ContextWrapper[T any] struct {
  92. pool ContextPool[T]
  93. }
  94. // NewContextWrapper returns a new ContextWrapper.
  95. // If pool is nil, a default pool is used.
  96. // The default pool's AcquireFunc returns a zero value of T.
  97. // The default pool's ReleaseFunc does nothing.
  98. // The default pool is used when the pool is nil.
  99. // Use the `iris.NewContextPool[T, *T]()` to pass a simple context pool.
  100. // Then, use the `Handler` method to wrap custom handlers to iris ones.
  101. //
  102. // Example: https://github.com/kataras/iris/tree/main/_examples/routing/custom-context
  103. func NewContextWrapper[T any](pool ContextPool[T]) *ContextWrapper[T] {
  104. if pool == nil {
  105. panic("pool cannot be nil")
  106. }
  107. return &ContextWrapper[T]{
  108. pool: pool,
  109. }
  110. }
  111. // Pool returns the pool, useful when manually Acquire and Release of custom context is required.
  112. func (w *ContextWrapper[T]) Pool() ContextPool[T] {
  113. return w.pool
  114. }
  115. // Handler wraps the handler with the pool's Acquire and Release methods.
  116. // It returns a new handler which expects a T instead of iris.Context.
  117. // The T is the type of the pool.
  118. // The T is acquired from the pool and released back to the pool after the handler's execution.
  119. // The T is passed to the handler as an argument.
  120. // The T is not shared between requests.
  121. func (w *ContextWrapper[T]) Handler(handler func(T)) Handler {
  122. if handler == nil {
  123. return nil
  124. }
  125. return func(ctx Context) {
  126. newT := w.pool.Acquire(ctx)
  127. handler(newT)
  128. w.pool.Release(newT)
  129. }
  130. }
  131. // Handlers wraps the handlers with the pool's Acquire and Release methods.
  132. func (w *ContextWrapper[T]) Handlers(handlers ...func(T)) context.Handlers {
  133. newHandlers := make(context.Handlers, len(handlers))
  134. for i, handler := range handlers {
  135. newHandlers[i] = w.Handler(handler)
  136. }
  137. return newHandlers
  138. }
  139. // HandlerReturnError same as `Handler` but it converts a handler which returns an error.
  140. func (w *ContextWrapper[T]) HandlerReturnError(handler func(T) error) func(Context) error {
  141. if handler == nil {
  142. return nil
  143. }
  144. return func(ctx Context) error {
  145. newT := w.pool.Acquire(ctx)
  146. err := handler(newT)
  147. w.pool.Release(newT)
  148. return err
  149. }
  150. }
  151. // HandlerReturnDuration same as `Handler` but it converts a handler which returns a time.Duration.
  152. func (w *ContextWrapper[T]) HandlerReturnDuration(handler func(T) time.Duration) func(Context) time.Duration {
  153. if handler == nil {
  154. return nil
  155. }
  156. return func(ctx Context) time.Duration {
  157. newT := w.pool.Acquire(ctx)
  158. duration := handler(newT)
  159. w.pool.Release(newT)
  160. return duration
  161. }
  162. }
  163. // Filter same as `Handler` but it converts a handler to Filter.
  164. func (w *ContextWrapper[T]) Filter(handler func(T) bool) Filter {
  165. if handler == nil {
  166. return nil
  167. }
  168. return func(ctx Context) bool {
  169. newT := w.pool.Acquire(ctx)
  170. shouldContinue := handler(newT)
  171. w.pool.Release(newT)
  172. return shouldContinue
  173. }
  174. }
  175. // FallbackViewFunc same as `Handler` but it converts a handler to FallbackViewFunc.
  176. func (w *ContextWrapper[T]) FallbackViewFunc(handler func(ctx T, err ErrViewNotExist) error) FallbackViewFunc {
  177. if handler == nil {
  178. return nil
  179. }
  180. return func(ctx Context, err ErrViewNotExist) error {
  181. newT := w.pool.Acquire(ctx)
  182. returningErr := handler(newT, err)
  183. w.pool.Release(newT)
  184. return returningErr
  185. }
  186. }