enhanced_memory_pools.go 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370
  1. package utils
  2. import (
  3. "strings"
  4. "sync"
  5. "time"
  6. )
  7. // EnhancedObjectPool provides advanced object pooling with automatic cleanup and monitoring
  8. type EnhancedObjectPool[T any] struct {
  9. pool sync.Pool
  10. created int64
  11. reused int64
  12. lastCleanup time.Time
  13. maxSize int
  14. resetFunc func(*T)
  15. mutex sync.RWMutex
  16. }
  17. // NewEnhancedObjectPool creates a new enhanced object pool
  18. func NewEnhancedObjectPool[T any](newFunc func() *T, resetFunc func(*T), maxSize int) *EnhancedObjectPool[T] {
  19. return &EnhancedObjectPool[T]{
  20. pool: sync.Pool{
  21. New: func() interface{} {
  22. return newFunc()
  23. },
  24. },
  25. maxSize: maxSize,
  26. resetFunc: resetFunc,
  27. lastCleanup: time.Now(),
  28. }
  29. }
  30. // Get retrieves an object from the pool
  31. func (p *EnhancedObjectPool[T]) Get() *T {
  32. obj := p.pool.Get().(*T)
  33. p.mutex.Lock()
  34. p.reused++
  35. p.mutex.Unlock()
  36. return obj
  37. }
  38. // Put returns an object to the pool after resetting it
  39. func (p *EnhancedObjectPool[T]) Put(obj *T) {
  40. if obj == nil {
  41. return
  42. }
  43. // Reset the object if reset function is provided
  44. if p.resetFunc != nil {
  45. p.resetFunc(obj)
  46. }
  47. p.pool.Put(obj)
  48. }
  49. // Stats returns pool statistics
  50. func (p *EnhancedObjectPool[T]) Stats() PoolStats {
  51. p.mutex.RLock()
  52. defer p.mutex.RUnlock()
  53. return PoolStats{
  54. Created: p.created,
  55. Reused: p.reused,
  56. ReuseRate: float64(p.reused) / float64(p.created+p.reused),
  57. LastAccess: p.lastCleanup,
  58. }
  59. }
  60. // PoolStats contains statistics about pool usage
  61. type PoolStats struct {
  62. Created int64 `json:"created"`
  63. Reused int64 `json:"reused"`
  64. ReuseRate float64 `json:"reuse_rate"`
  65. LastAccess time.Time `json:"last_access"`
  66. }
  67. // StringBuilderPool provides pooled string builders
  68. type StringBuilderPool struct {
  69. pool *EnhancedObjectPool[strings.Builder]
  70. }
  71. // NewStringBuilderPool creates a new string builder pool
  72. func NewStringBuilderPool(initialCap, maxSize int) *StringBuilderPool {
  73. return &StringBuilderPool{
  74. pool: NewEnhancedObjectPool(
  75. func() *strings.Builder {
  76. sb := &strings.Builder{}
  77. sb.Grow(initialCap)
  78. return sb
  79. },
  80. func(sb *strings.Builder) { sb.Reset() },
  81. maxSize,
  82. ),
  83. }
  84. }
  85. // Get retrieves a string builder from the pool
  86. func (p *StringBuilderPool) Get() *strings.Builder {
  87. return p.pool.Get()
  88. }
  89. // Put returns a string builder to the pool
  90. func (p *StringBuilderPool) Put(sb *strings.Builder) {
  91. p.pool.Put(sb)
  92. }
  93. // ByteSlicePool provides pooled byte slices
  94. type ByteSlicePool struct {
  95. pools map[int]*EnhancedObjectPool[[]byte]
  96. mutex sync.RWMutex
  97. }
  98. // NewByteSlicePool creates a new byte slice pool
  99. func NewByteSlicePool() *ByteSlicePool {
  100. return &ByteSlicePool{
  101. pools: make(map[int]*EnhancedObjectPool[[]byte]),
  102. }
  103. }
  104. // Get retrieves a byte slice of the requested size
  105. func (p *ByteSlicePool) Get(size int) []byte {
  106. // Round up to nearest power of 2 for better pooling
  107. poolSize := nextPowerOf2(size)
  108. p.mutex.RLock()
  109. pool, exists := p.pools[poolSize]
  110. p.mutex.RUnlock()
  111. if !exists {
  112. p.mutex.Lock()
  113. // Double-check after acquiring write lock
  114. if pool, exists = p.pools[poolSize]; !exists {
  115. pool = NewEnhancedObjectPool(
  116. func() *[]byte {
  117. slice := make([]byte, 0, poolSize)
  118. return &slice
  119. },
  120. func(slice *[]byte) { *slice = (*slice)[:0] },
  121. 100, // max 100 slices per size
  122. )
  123. p.pools[poolSize] = pool
  124. }
  125. p.mutex.Unlock()
  126. }
  127. slice := pool.Get()
  128. return *slice
  129. }
  130. // Put returns a byte slice to the pool
  131. func (p *ByteSlicePool) Put(slice []byte) {
  132. if slice == nil {
  133. return
  134. }
  135. capacity := cap(slice)
  136. p.mutex.RLock()
  137. pool, exists := p.pools[capacity]
  138. p.mutex.RUnlock()
  139. if exists {
  140. pool.Put(&slice)
  141. }
  142. }
  143. // nextPowerOf2 returns the next power of 2 greater than or equal to n
  144. func nextPowerOf2(n int) int {
  145. if n <= 0 {
  146. return 1
  147. }
  148. n--
  149. n |= n >> 1
  150. n |= n >> 2
  151. n |= n >> 4
  152. n |= n >> 8
  153. n |= n >> 16
  154. n++
  155. return n
  156. }
  157. // MapPool provides pooled maps
  158. type MapPool[K comparable, V any] struct {
  159. pool *EnhancedObjectPool[map[K]V]
  160. }
  161. // NewMapPool creates a new map pool
  162. func NewMapPool[K comparable, V any](initialSize, maxSize int) *MapPool[K, V] {
  163. return &MapPool[K, V]{
  164. pool: NewEnhancedObjectPool(
  165. func() *map[K]V {
  166. m := make(map[K]V, initialSize)
  167. return &m
  168. },
  169. func(m *map[K]V) {
  170. // Clear the map
  171. for k := range *m {
  172. delete(*m, k)
  173. }
  174. },
  175. maxSize,
  176. ),
  177. }
  178. }
  179. // Get retrieves a map from the pool
  180. func (p *MapPool[K, V]) Get() map[K]V {
  181. return *p.pool.Get()
  182. }
  183. // Put returns a map to the pool
  184. func (p *MapPool[K, V]) Put(m map[K]V) {
  185. p.pool.Put(&m)
  186. }
  187. // SlicePool provides pooled slices
  188. type SlicePool[T any] struct {
  189. pool *EnhancedObjectPool[[]T]
  190. }
  191. // NewSlicePool creates a new slice pool
  192. func NewSlicePool[T any](initialCap, maxSize int) *SlicePool[T] {
  193. return &SlicePool[T]{
  194. pool: NewEnhancedObjectPool(
  195. func() *[]T {
  196. slice := make([]T, 0, initialCap)
  197. return &slice
  198. },
  199. func(slice *[]T) { *slice = (*slice)[:0] },
  200. maxSize,
  201. ),
  202. }
  203. }
  204. // Get retrieves a slice from the pool
  205. func (p *SlicePool[T]) Get() []T {
  206. return *p.pool.Get()
  207. }
  208. // Put returns a slice to the pool
  209. func (p *SlicePool[T]) Put(slice []T) {
  210. p.pool.Put(&slice)
  211. }
  212. // PoolManager manages multiple object pools
  213. type PoolManager struct {
  214. pools map[string]interface{}
  215. mutex sync.RWMutex
  216. }
  217. // NewPoolManager creates a new pool manager
  218. func NewPoolManager() *PoolManager {
  219. return &PoolManager{
  220. pools: make(map[string]interface{}),
  221. }
  222. }
  223. // RegisterPool registers a pool with the manager
  224. func (pm *PoolManager) RegisterPool(name string, pool interface{}) {
  225. pm.mutex.Lock()
  226. defer pm.mutex.Unlock()
  227. pm.pools[name] = pool
  228. }
  229. // GetPool retrieves a pool by name
  230. func (pm *PoolManager) GetPool(name string) (interface{}, bool) {
  231. pm.mutex.RLock()
  232. defer pm.mutex.RUnlock()
  233. pool, exists := pm.pools[name]
  234. return pool, exists
  235. }
  236. // GetAllStats returns statistics for all registered pools
  237. func (pm *PoolManager) GetAllStats() map[string]interface{} {
  238. pm.mutex.RLock()
  239. defer pm.mutex.RUnlock()
  240. stats := make(map[string]interface{})
  241. for name, pool := range pm.pools {
  242. // Try to get stats if the pool supports it
  243. if statsProvider, ok := pool.(interface{ Stats() PoolStats }); ok {
  244. stats[name] = statsProvider.Stats()
  245. } else {
  246. stats[name] = "stats not available"
  247. }
  248. }
  249. return stats
  250. }
  251. // Global pool manager instance
  252. var globalPoolManager = NewPoolManager()
  253. // GetGlobalPoolManager returns the global pool manager
  254. func GetGlobalPoolManager() *PoolManager {
  255. return globalPoolManager
  256. }
  257. // Common pool instances for frequent use
  258. var (
  259. // String builder pool for log processing
  260. LogStringBuilderPool = NewStringBuilderPool(1024, 50)
  261. // Byte slice pool for I/O operations
  262. GlobalByteSlicePool = NewByteSlicePool()
  263. // String slice pool for batch processing
  264. StringSlicePool = NewSlicePool[string](100, 20)
  265. // Map pools for common use cases
  266. StringMapPool = NewMapPool[string, string](10, 20)
  267. StringIntMapPool = NewMapPool[string, int](10, 20)
  268. IntStringMapPool = NewMapPool[int, string](10, 20)
  269. )
  270. // Initialize global pools
  271. func init() {
  272. // Register common pools with the global manager
  273. globalPoolManager.RegisterPool("log_string_builder", LogStringBuilderPool)
  274. globalPoolManager.RegisterPool("global_byte_slice", GlobalByteSlicePool)
  275. globalPoolManager.RegisterPool("string_slice", StringSlicePool)
  276. globalPoolManager.RegisterPool("string_map", StringMapPool)
  277. globalPoolManager.RegisterPool("string_int_map", StringIntMapPool)
  278. globalPoolManager.RegisterPool("int_string_map", IntStringMapPool)
  279. }
  280. // PooledWorker represents a worker that uses object pools
  281. type PooledWorker struct {
  282. stringBuilders *StringBuilderPool
  283. byteSlices *ByteSlicePool
  284. workBuffer []byte
  285. }
  286. // NewPooledWorker creates a new pooled worker
  287. func NewPooledWorker() *PooledWorker {
  288. return &PooledWorker{
  289. stringBuilders: LogStringBuilderPool,
  290. byteSlices: GlobalByteSlicePool,
  291. }
  292. }
  293. // ProcessWithPools processes data using object pools to minimize allocations
  294. func (pw *PooledWorker) ProcessWithPools(data []byte, processor func([]byte, *strings.Builder) error) error {
  295. // Get pooled string builder
  296. sb := pw.stringBuilders.Get()
  297. defer pw.stringBuilders.Put(sb)
  298. // Get pooled byte slice if needed
  299. if len(pw.workBuffer) < len(data) {
  300. if pw.workBuffer != nil {
  301. pw.byteSlices.Put(pw.workBuffer)
  302. }
  303. pw.workBuffer = pw.byteSlices.Get(len(data))
  304. }
  305. // Use pooled objects for processing
  306. copy(pw.workBuffer, data)
  307. return processor(pw.workBuffer, sb)
  308. }
  309. // Cleanup releases resources held by the worker
  310. func (pw *PooledWorker) Cleanup() {
  311. if pw.workBuffer != nil {
  312. pw.byteSlices.Put(pw.workBuffer)
  313. pw.workBuffer = nil
  314. }
  315. }