| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319 |
- package fiber
- import (
- "encoding/hex"
- "strings"
- "sync"
- "github.com/google/uuid"
- )
- const servicesStatePrefix = "gofiber-services-"
- var servicesStatePrefixHash string
- func init() {
- servicesStatePrefixHash = hex.EncodeToString([]byte(servicesStatePrefix + uuid.New().String()))
- }
- // State is a key-value store for Fiber's app in order to be used as a global storage for the app's dependencies.
- // It's a thread-safe implementation of a map[string]any, using sync.Map.
- type State struct {
- dependencies sync.Map
- servicePrefix string
- }
- // NewState creates a new instance of State.
- func newState() *State {
- // Initialize the services state prefix using a hashed random string
- return &State{
- dependencies: sync.Map{},
- servicePrefix: servicesStatePrefixHash,
- }
- }
- // Set sets a key-value pair in the State.
- func (s *State) Set(key string, value any) {
- s.dependencies.Store(key, value)
- }
- // Get retrieves a value from the State.
- func (s *State) Get(key string) (any, bool) {
- return s.dependencies.Load(key)
- }
- // MustGet retrieves a value from the State and panics if the key is not found.
- func (s *State) MustGet(key string) any {
- if dep, ok := s.Get(key); ok {
- return dep
- }
- panic("state: dependency not found!")
- }
- // Has checks if a key is present in the State.
- // It returns a boolean indicating if the key is present.
- func (s *State) Has(key string) bool {
- _, ok := s.Get(key)
- return ok
- }
- // Delete removes a key-value pair from the State.
- func (s *State) Delete(key string) {
- s.dependencies.Delete(key)
- }
- // Reset resets the State by removing all keys.
- func (s *State) Reset() {
- s.dependencies.Clear()
- }
- // Keys returns a slice containing all keys present in the State.
- func (s *State) Keys() []string {
- // Pre-allocate with estimated capacity to reduce allocations
- keys := make([]string, 0, 8)
- s.dependencies.Range(func(key, _ any) bool {
- keyStr, ok := key.(string)
- if !ok {
- return true
- }
- keys = append(keys, keyStr)
- return true
- })
- return keys
- }
- // Len returns the number of keys in the State.
- func (s *State) Len() int {
- length := 0
- s.dependencies.Range(func(_, _ any) bool {
- length++
- return true
- })
- return length
- }
- // GetState retrieves a value from the State and casts it to the desired type.
- // It returns the casted value and a boolean indicating if the cast was successful.
- func GetState[T any](s *State, key string) (T, bool) {
- dep, ok := s.Get(key)
- if ok {
- depT, okCast := dep.(T)
- return depT, okCast
- }
- var zeroVal T
- return zeroVal, false
- }
- // MustGetState retrieves a value from the State and casts it to the desired type.
- // It panics if the key is not found or if the type assertion fails.
- func MustGetState[T any](s *State, key string) T {
- dep, ok := GetState[T](s, key)
- if !ok {
- panic("state: dependency not found!")
- }
- return dep
- }
- // GetStateWithDefault retrieves a value from the State,
- // casting it to the desired type. If the key is not present,
- // it returns the provided default value.
- func GetStateWithDefault[T any](s *State, key string, defaultVal T) T {
- dep, ok := GetState[T](s, key)
- if !ok {
- return defaultVal
- }
- return dep
- }
- // GetString retrieves a string value from the State.
- // It returns the string and a boolean indicating successful type assertion.
- func (s *State) GetString(key string) (string, bool) {
- return GetState[string](s, key)
- }
- // GetInt retrieves an integer value from the State.
- // It returns the int and a boolean indicating successful type assertion.
- func (s *State) GetInt(key string) (int, bool) {
- return GetState[int](s, key)
- }
- // GetBool retrieves a boolean value from the State.
- // It returns the bool and a boolean indicating successful type assertion.
- func (s *State) GetBool(key string) (value, ok bool) { //nolint:nonamedreturns // Better idea to use named returns here
- return GetState[bool](s, key)
- }
- // GetFloat64 retrieves a float64 value from the State.
- // It returns the float64 and a boolean indicating successful type assertion.
- func (s *State) GetFloat64(key string) (float64, bool) {
- return GetState[float64](s, key)
- }
- // GetUint retrieves a uint value from the State.
- // It returns the uint and a boolean indicating successful type assertion.
- func (s *State) GetUint(key string) (uint, bool) {
- return GetState[uint](s, key)
- }
- // GetInt8 retrieves an int8 value from the State.
- // It returns the int8 and a boolean indicating successful type assertion.
- func (s *State) GetInt8(key string) (int8, bool) {
- return GetState[int8](s, key)
- }
- // GetInt16 retrieves an int16 value from the State.
- // It returns the int16 and a boolean indicating successful type assertion.
- func (s *State) GetInt16(key string) (int16, bool) {
- return GetState[int16](s, key)
- }
- // GetInt32 retrieves an int32 value from the State.
- // It returns the int32 and a boolean indicating successful type assertion.
- func (s *State) GetInt32(key string) (int32, bool) {
- return GetState[int32](s, key)
- }
- // GetInt64 retrieves an int64 value from the State.
- // It returns the int64 and a boolean indicating successful type assertion.
- func (s *State) GetInt64(key string) (int64, bool) {
- return GetState[int64](s, key)
- }
- // GetUint8 retrieves a uint8 value from the State.
- // It returns the uint8 and a boolean indicating successful type assertion.
- func (s *State) GetUint8(key string) (uint8, bool) {
- return GetState[uint8](s, key)
- }
- // GetUint16 retrieves a uint16 value from the State.
- // It returns the uint16 and a boolean indicating successful type assertion.
- func (s *State) GetUint16(key string) (uint16, bool) {
- return GetState[uint16](s, key)
- }
- // GetUint32 retrieves a uint32 value from the State.
- // It returns the uint32 and a boolean indicating successful type assertion.
- func (s *State) GetUint32(key string) (uint32, bool) {
- return GetState[uint32](s, key)
- }
- // GetUint64 retrieves a uint64 value from the State.
- // It returns the uint64 and a boolean indicating successful type assertion.
- func (s *State) GetUint64(key string) (uint64, bool) {
- return GetState[uint64](s, key)
- }
- // GetUintptr retrieves a uintptr value from the State.
- // It returns the uintptr and a boolean indicating successful type assertion.
- func (s *State) GetUintptr(key string) (uintptr, bool) {
- return GetState[uintptr](s, key)
- }
- // GetFloat32 retrieves a float32 value from the State.
- // It returns the float32 and a boolean indicating successful type assertion.
- func (s *State) GetFloat32(key string) (float32, bool) {
- return GetState[float32](s, key)
- }
- // GetComplex64 retrieves a complex64 value from the State.
- // It returns the complex64 and a boolean indicating successful type assertion.
- func (s *State) GetComplex64(key string) (complex64, bool) {
- return GetState[complex64](s, key)
- }
- // GetComplex128 retrieves a complex128 value from the State.
- // It returns the complex128 and a boolean indicating successful type assertion.
- func (s *State) GetComplex128(key string) (complex128, bool) {
- return GetState[complex128](s, key)
- }
- // serviceKey returns a key for a service in the State.
- // A key is composed of the State's servicePrefix (hashed) and the hash of the service string.
- // This way we can avoid collisions and have a unique key for each service.
- func (s *State) serviceKey(key string) string {
- // hash the service string to avoid collisions
- return s.servicePrefix + hex.EncodeToString([]byte(key))
- }
- // setService sets a service in the State.
- func (s *State) setService(srv Service) {
- // Always prepend the service key with the servicesStateKey to avoid collisions
- s.Set(s.serviceKey(srv.String()), srv)
- }
- // Delete removes a key-value pair from the State.
- func (s *State) deleteService(srv Service) {
- s.Delete(s.serviceKey(srv.String()))
- }
- // serviceKeys returns a slice containing all keys present for services in the application's State.
- func (s *State) serviceKeys() []string {
- // Pre-allocate with estimated capacity to reduce allocations
- keys := make([]string, 0, 8)
- s.dependencies.Range(func(key, _ any) bool {
- keyStr, ok := key.(string)
- if !ok {
- return true
- }
- if !strings.HasPrefix(keyStr, s.servicePrefix) {
- return true // Continue iterating if key doesn't have service prefix
- }
- keys = append(keys, keyStr)
- return true
- })
- return keys
- }
- // Services returns a map containing all services present in the State.
- // The key is the hash of the service String() value and the value is the service itself.
- func (s *State) Services() map[string]Service {
- keys := s.serviceKeys()
- services := make(map[string]Service, len(keys))
- for _, key := range keys {
- services[key] = MustGetState[Service](s, key)
- }
- return services
- }
- // ServicesLen returns the number of keys for services in the State.
- func (s *State) ServicesLen() int {
- length := 0
- s.dependencies.Range(func(key, _ any) bool {
- if str, ok := key.(string); ok && strings.HasPrefix(str, s.servicePrefix) {
- length++
- }
- return true
- })
- return length
- }
- // GetService returns a service present in the application's State.
- func GetService[T Service](s *State, key string) (T, bool) {
- srv, ok := GetState[T](s, s.serviceKey(key))
- return srv, ok
- }
- // MustGetService returns a service present in the application's State.
- // It panics if the service is not found.
- func MustGetService[T Service](s *State, key string) T {
- srv, ok := GetService[T](s, key)
- if !ok {
- panic("state: service not found!")
- }
- return srv
- }
|