state.go 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319
  1. package fiber
  2. import (
  3. "encoding/hex"
  4. "strings"
  5. "sync"
  6. "github.com/google/uuid"
  7. )
  8. const servicesStatePrefix = "gofiber-services-"
  9. var servicesStatePrefixHash string
  10. func init() {
  11. servicesStatePrefixHash = hex.EncodeToString([]byte(servicesStatePrefix + uuid.New().String()))
  12. }
  13. // State is a key-value store for Fiber's app in order to be used as a global storage for the app's dependencies.
  14. // It's a thread-safe implementation of a map[string]any, using sync.Map.
  15. type State struct {
  16. dependencies sync.Map
  17. servicePrefix string
  18. }
  19. // NewState creates a new instance of State.
  20. func newState() *State {
  21. // Initialize the services state prefix using a hashed random string
  22. return &State{
  23. dependencies: sync.Map{},
  24. servicePrefix: servicesStatePrefixHash,
  25. }
  26. }
  27. // Set sets a key-value pair in the State.
  28. func (s *State) Set(key string, value any) {
  29. s.dependencies.Store(key, value)
  30. }
  31. // Get retrieves a value from the State.
  32. func (s *State) Get(key string) (any, bool) {
  33. return s.dependencies.Load(key)
  34. }
  35. // MustGet retrieves a value from the State and panics if the key is not found.
  36. func (s *State) MustGet(key string) any {
  37. if dep, ok := s.Get(key); ok {
  38. return dep
  39. }
  40. panic("state: dependency not found!")
  41. }
  42. // Has checks if a key is present in the State.
  43. // It returns a boolean indicating if the key is present.
  44. func (s *State) Has(key string) bool {
  45. _, ok := s.Get(key)
  46. return ok
  47. }
  48. // Delete removes a key-value pair from the State.
  49. func (s *State) Delete(key string) {
  50. s.dependencies.Delete(key)
  51. }
  52. // Reset resets the State by removing all keys.
  53. func (s *State) Reset() {
  54. s.dependencies.Clear()
  55. }
  56. // Keys returns a slice containing all keys present in the State.
  57. func (s *State) Keys() []string {
  58. // Pre-allocate with estimated capacity to reduce allocations
  59. keys := make([]string, 0, 8)
  60. s.dependencies.Range(func(key, _ any) bool {
  61. keyStr, ok := key.(string)
  62. if !ok {
  63. return true
  64. }
  65. keys = append(keys, keyStr)
  66. return true
  67. })
  68. return keys
  69. }
  70. // Len returns the number of keys in the State.
  71. func (s *State) Len() int {
  72. length := 0
  73. s.dependencies.Range(func(_, _ any) bool {
  74. length++
  75. return true
  76. })
  77. return length
  78. }
  79. // GetState retrieves a value from the State and casts it to the desired type.
  80. // It returns the casted value and a boolean indicating if the cast was successful.
  81. func GetState[T any](s *State, key string) (T, bool) {
  82. dep, ok := s.Get(key)
  83. if ok {
  84. depT, okCast := dep.(T)
  85. return depT, okCast
  86. }
  87. var zeroVal T
  88. return zeroVal, false
  89. }
  90. // MustGetState retrieves a value from the State and casts it to the desired type.
  91. // It panics if the key is not found or if the type assertion fails.
  92. func MustGetState[T any](s *State, key string) T {
  93. dep, ok := GetState[T](s, key)
  94. if !ok {
  95. panic("state: dependency not found!")
  96. }
  97. return dep
  98. }
  99. // GetStateWithDefault retrieves a value from the State,
  100. // casting it to the desired type. If the key is not present,
  101. // it returns the provided default value.
  102. func GetStateWithDefault[T any](s *State, key string, defaultVal T) T {
  103. dep, ok := GetState[T](s, key)
  104. if !ok {
  105. return defaultVal
  106. }
  107. return dep
  108. }
  109. // GetString retrieves a string value from the State.
  110. // It returns the string and a boolean indicating successful type assertion.
  111. func (s *State) GetString(key string) (string, bool) {
  112. return GetState[string](s, key)
  113. }
  114. // GetInt retrieves an integer value from the State.
  115. // It returns the int and a boolean indicating successful type assertion.
  116. func (s *State) GetInt(key string) (int, bool) {
  117. return GetState[int](s, key)
  118. }
  119. // GetBool retrieves a boolean value from the State.
  120. // It returns the bool and a boolean indicating successful type assertion.
  121. func (s *State) GetBool(key string) (value, ok bool) { //nolint:nonamedreturns // Better idea to use named returns here
  122. return GetState[bool](s, key)
  123. }
  124. // GetFloat64 retrieves a float64 value from the State.
  125. // It returns the float64 and a boolean indicating successful type assertion.
  126. func (s *State) GetFloat64(key string) (float64, bool) {
  127. return GetState[float64](s, key)
  128. }
  129. // GetUint retrieves a uint value from the State.
  130. // It returns the uint and a boolean indicating successful type assertion.
  131. func (s *State) GetUint(key string) (uint, bool) {
  132. return GetState[uint](s, key)
  133. }
  134. // GetInt8 retrieves an int8 value from the State.
  135. // It returns the int8 and a boolean indicating successful type assertion.
  136. func (s *State) GetInt8(key string) (int8, bool) {
  137. return GetState[int8](s, key)
  138. }
  139. // GetInt16 retrieves an int16 value from the State.
  140. // It returns the int16 and a boolean indicating successful type assertion.
  141. func (s *State) GetInt16(key string) (int16, bool) {
  142. return GetState[int16](s, key)
  143. }
  144. // GetInt32 retrieves an int32 value from the State.
  145. // It returns the int32 and a boolean indicating successful type assertion.
  146. func (s *State) GetInt32(key string) (int32, bool) {
  147. return GetState[int32](s, key)
  148. }
  149. // GetInt64 retrieves an int64 value from the State.
  150. // It returns the int64 and a boolean indicating successful type assertion.
  151. func (s *State) GetInt64(key string) (int64, bool) {
  152. return GetState[int64](s, key)
  153. }
  154. // GetUint8 retrieves a uint8 value from the State.
  155. // It returns the uint8 and a boolean indicating successful type assertion.
  156. func (s *State) GetUint8(key string) (uint8, bool) {
  157. return GetState[uint8](s, key)
  158. }
  159. // GetUint16 retrieves a uint16 value from the State.
  160. // It returns the uint16 and a boolean indicating successful type assertion.
  161. func (s *State) GetUint16(key string) (uint16, bool) {
  162. return GetState[uint16](s, key)
  163. }
  164. // GetUint32 retrieves a uint32 value from the State.
  165. // It returns the uint32 and a boolean indicating successful type assertion.
  166. func (s *State) GetUint32(key string) (uint32, bool) {
  167. return GetState[uint32](s, key)
  168. }
  169. // GetUint64 retrieves a uint64 value from the State.
  170. // It returns the uint64 and a boolean indicating successful type assertion.
  171. func (s *State) GetUint64(key string) (uint64, bool) {
  172. return GetState[uint64](s, key)
  173. }
  174. // GetUintptr retrieves a uintptr value from the State.
  175. // It returns the uintptr and a boolean indicating successful type assertion.
  176. func (s *State) GetUintptr(key string) (uintptr, bool) {
  177. return GetState[uintptr](s, key)
  178. }
  179. // GetFloat32 retrieves a float32 value from the State.
  180. // It returns the float32 and a boolean indicating successful type assertion.
  181. func (s *State) GetFloat32(key string) (float32, bool) {
  182. return GetState[float32](s, key)
  183. }
  184. // GetComplex64 retrieves a complex64 value from the State.
  185. // It returns the complex64 and a boolean indicating successful type assertion.
  186. func (s *State) GetComplex64(key string) (complex64, bool) {
  187. return GetState[complex64](s, key)
  188. }
  189. // GetComplex128 retrieves a complex128 value from the State.
  190. // It returns the complex128 and a boolean indicating successful type assertion.
  191. func (s *State) GetComplex128(key string) (complex128, bool) {
  192. return GetState[complex128](s, key)
  193. }
  194. // serviceKey returns a key for a service in the State.
  195. // A key is composed of the State's servicePrefix (hashed) and the hash of the service string.
  196. // This way we can avoid collisions and have a unique key for each service.
  197. func (s *State) serviceKey(key string) string {
  198. // hash the service string to avoid collisions
  199. return s.servicePrefix + hex.EncodeToString([]byte(key))
  200. }
  201. // setService sets a service in the State.
  202. func (s *State) setService(srv Service) {
  203. // Always prepend the service key with the servicesStateKey to avoid collisions
  204. s.Set(s.serviceKey(srv.String()), srv)
  205. }
  206. // Delete removes a key-value pair from the State.
  207. func (s *State) deleteService(srv Service) {
  208. s.Delete(s.serviceKey(srv.String()))
  209. }
  210. // serviceKeys returns a slice containing all keys present for services in the application's State.
  211. func (s *State) serviceKeys() []string {
  212. // Pre-allocate with estimated capacity to reduce allocations
  213. keys := make([]string, 0, 8)
  214. s.dependencies.Range(func(key, _ any) bool {
  215. keyStr, ok := key.(string)
  216. if !ok {
  217. return true
  218. }
  219. if !strings.HasPrefix(keyStr, s.servicePrefix) {
  220. return true // Continue iterating if key doesn't have service prefix
  221. }
  222. keys = append(keys, keyStr)
  223. return true
  224. })
  225. return keys
  226. }
  227. // Services returns a map containing all services present in the State.
  228. // The key is the hash of the service String() value and the value is the service itself.
  229. func (s *State) Services() map[string]Service {
  230. keys := s.serviceKeys()
  231. services := make(map[string]Service, len(keys))
  232. for _, key := range keys {
  233. services[key] = MustGetState[Service](s, key)
  234. }
  235. return services
  236. }
  237. // ServicesLen returns the number of keys for services in the State.
  238. func (s *State) ServicesLen() int {
  239. length := 0
  240. s.dependencies.Range(func(key, _ any) bool {
  241. if str, ok := key.(string); ok && strings.HasPrefix(str, s.servicePrefix) {
  242. length++
  243. }
  244. return true
  245. })
  246. return length
  247. }
  248. // GetService returns a service present in the application's State.
  249. func GetService[T Service](s *State, key string) (T, bool) {
  250. srv, ok := GetState[T](s, s.serviceKey(key))
  251. return srv, ok
  252. }
  253. // MustGetService returns a service present in the application's State.
  254. // It panics if the service is not found.
  255. func MustGetService[T Service](s *State, key string) T {
  256. srv, ok := GetService[T](s, key)
  257. if !ok {
  258. panic("state: service not found!")
  259. }
  260. return srv
  261. }