hooks.go 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. package fiber
  2. // OnRouteHandler Handlers define a function to create hooks for Fiber.
  3. type OnRouteHandler = func(Route) error
  4. type OnNameHandler = OnRouteHandler
  5. type OnGroupHandler = func(Group) error
  6. type OnGroupNameHandler = OnGroupHandler
  7. type OnListenHandler = func() error
  8. type OnShutdownHandler = OnListenHandler
  9. type OnForkHandler = func(int) error
  10. // Hooks is a struct to use it with App.
  11. type Hooks struct {
  12. // Embed app
  13. app *App
  14. // Hooks
  15. onRoute []OnRouteHandler
  16. onName []OnNameHandler
  17. onGroup []OnGroupHandler
  18. onGroupName []OnGroupNameHandler
  19. onListen []OnListenHandler
  20. onShutdown []OnShutdownHandler
  21. onFork []OnForkHandler
  22. }
  23. func newHooks(app *App) *Hooks {
  24. return &Hooks{
  25. app: app,
  26. onRoute: make([]OnRouteHandler, 0),
  27. onGroup: make([]OnGroupHandler, 0),
  28. onGroupName: make([]OnGroupNameHandler, 0),
  29. onName: make([]OnNameHandler, 0),
  30. onListen: make([]OnListenHandler, 0),
  31. onShutdown: make([]OnShutdownHandler, 0),
  32. onFork: make([]OnForkHandler, 0),
  33. }
  34. }
  35. // OnRoute is a hook to execute user functions on each route registeration.
  36. // Also you can get route properties by route parameter.
  37. func (h *Hooks) OnRoute(handler ...OnRouteHandler) {
  38. h.app.mutex.Lock()
  39. h.onRoute = append(h.onRoute, handler...)
  40. h.app.mutex.Unlock()
  41. }
  42. // OnName is a hook to execute user functions on each route naming.
  43. // Also you can get route properties by route parameter.
  44. //
  45. // WARN: OnName only works with naming routes, not groups.
  46. func (h *Hooks) OnName(handler ...OnNameHandler) {
  47. h.app.mutex.Lock()
  48. h.onName = append(h.onName, handler...)
  49. h.app.mutex.Unlock()
  50. }
  51. // OnGroup is a hook to execute user functions on each group registeration.
  52. // Also you can get group properties by group parameter.
  53. func (h *Hooks) OnGroup(handler ...OnGroupHandler) {
  54. h.app.mutex.Lock()
  55. h.onGroup = append(h.onGroup, handler...)
  56. h.app.mutex.Unlock()
  57. }
  58. // OnGroupName is a hook to execute user functions on each group naming.
  59. // Also you can get group properties by group parameter.
  60. //
  61. // WARN: OnGroupName only works with naming groups, not routes.
  62. func (h *Hooks) OnGroupName(handler ...OnGroupNameHandler) {
  63. h.app.mutex.Lock()
  64. h.onGroupName = append(h.onGroupName, handler...)
  65. h.app.mutex.Unlock()
  66. }
  67. // OnListen is a hook to execute user functions on Listen, ListenTLS, Listener.
  68. func (h *Hooks) OnListen(handler ...OnListenHandler) {
  69. h.app.mutex.Lock()
  70. h.onListen = append(h.onListen, handler...)
  71. h.app.mutex.Unlock()
  72. }
  73. // OnShutdown is a hook to execute user functions after Shutdown.
  74. func (h *Hooks) OnShutdown(handler ...OnShutdownHandler) {
  75. h.app.mutex.Lock()
  76. h.onShutdown = append(h.onShutdown, handler...)
  77. h.app.mutex.Unlock()
  78. }
  79. // OnFork is a hook to execute user function after fork process.
  80. func (h *Hooks) OnFork(handler ...OnForkHandler) {
  81. h.app.mutex.Lock()
  82. h.onFork = append(h.onFork, handler...)
  83. h.app.mutex.Unlock()
  84. }
  85. func (h *Hooks) executeOnRouteHooks(route Route) error {
  86. for _, v := range h.onRoute {
  87. if err := v(route); err != nil {
  88. return err
  89. }
  90. }
  91. return nil
  92. }
  93. func (h *Hooks) executeOnNameHooks(route Route) error {
  94. for _, v := range h.onName {
  95. if err := v(route); err != nil {
  96. return err
  97. }
  98. }
  99. return nil
  100. }
  101. func (h *Hooks) executeOnGroupHooks(group Group) error {
  102. for _, v := range h.onGroup {
  103. if err := v(group); err != nil {
  104. return err
  105. }
  106. }
  107. return nil
  108. }
  109. func (h *Hooks) executeOnGroupNameHooks(group Group) error {
  110. for _, v := range h.onGroupName {
  111. if err := v(group); err != nil {
  112. return err
  113. }
  114. }
  115. return nil
  116. }
  117. func (h *Hooks) executeOnListenHooks() error {
  118. for _, v := range h.onListen {
  119. if err := v(); err != nil {
  120. return err
  121. }
  122. }
  123. return nil
  124. }
  125. func (h *Hooks) executeOnShutdownHooks() {
  126. for _, v := range h.onShutdown {
  127. _ = v()
  128. }
  129. }
  130. func (h *Hooks) executeOnForkHooks(pid int) {
  131. for _, v := range h.onFork {
  132. _ = v(pid)
  133. }
  134. }