@@ -147,14 +147,10 @@ func New[PT Item[T], T any](
147147 opts ... option [PT , T ],
148148) * Pool [PT , T ] {
149149 p := & Pool [PT , T ]{
150- trace : defaultTrace ,
151- limit : DefaultLimit ,
152- createItem : func (ctx context.Context ) (PT , error ) {
153- var item T
154-
155- return & item , nil
156- },
157- done : make (chan struct {}),
150+ trace : defaultTrace ,
151+ limit : DefaultLimit ,
152+ createItem : defaultCreateItem [T , PT ],
153+ done : make (chan struct {}),
158154 }
159155
160156 for _ , opt := range opts {
@@ -174,62 +170,38 @@ func New[PT Item[T], T any](
174170 })
175171 }()
176172
177- createItem := p .createItem
173+ p .createItem = createItemWithTimeoutHandling (p .createItem , p )
174+
175+ p .idle = make ([]PT , 0 , p .limit )
176+ p .index = make (map [PT ]struct {}, p .limit )
177+ p .stats = & safeStats {
178+ v : stats.Stats {Limit : p .limit },
179+ onChange : p .trace .OnChange ,
180+ }
181+
182+ return p
183+ }
184+
185+ // defaultCreateItem returns a new item
186+ func defaultCreateItem [T any , PT Item [T ]](ctx context.Context ) (PT , error ) {
187+ var item T
188+
189+ return & item , nil
190+ }
178191
179- p .createItem = func (ctx context.Context ) (PT , error ) {
192+ // createItemWithTimeoutHandling wraps the createItem function with timeout handling
193+ func createItemWithTimeoutHandling [PT Item [T ], T any ](
194+ createItem func (ctx context.Context ) (PT , error ),
195+ p * Pool [PT , T ],
196+ ) func (ctx context.Context ) (PT , error ) {
197+ return func (ctx context.Context ) (PT , error ) {
180198 var (
181199 ch = make (chan PT )
182200 createErr error
183201 )
184202 go func () {
185203 defer close (ch )
186- createErr = func () error {
187- var (
188- createCtx = xcontext .ValueOnly (ctx )
189- cancelCreate context.CancelFunc
190- )
191- if d := p .createTimeout ; d > 0 {
192- createCtx , cancelCreate = xcontext .WithTimeout (createCtx , d )
193- } else {
194- createCtx , cancelCreate = xcontext .WithCancel (createCtx )
195- }
196- defer cancelCreate ()
197-
198- newItem , err := createItem (createCtx )
199- if err != nil {
200- return xerrors .WithStackTrace (err )
201- }
202-
203- needCloseItem := true
204- defer func () {
205- if needCloseItem {
206- _ = p .closeItem (ctx , newItem )
207- }
208- }()
209-
210- select {
211- case <- p .done :
212- return xerrors .WithStackTrace (errClosedPool )
213-
214- case <- ctx .Done ():
215- p .mu .Lock ()
216- defer p .mu .Unlock ()
217-
218- if len (p .index ) < p .limit {
219- p .idle = append (p .idle , newItem )
220- p .index [newItem ] = struct {}{}
221- p .stats .Index ().Inc ()
222- needCloseItem = false
223- }
224-
225- return xerrors .WithStackTrace (ctx .Err ())
226-
227- case ch <- newItem :
228- needCloseItem = false
229-
230- return nil
231- }
232- }()
204+ createErr = createItemWithContext (ctx , p , createItem , ch )
233205 }()
234206
235207 select {
@@ -249,14 +221,59 @@ func New[PT Item[T], T any](
249221 return item , nil
250222 }
251223 }
252- p .idle = make ([]PT , 0 , p .limit )
253- p .index = make (map [PT ]struct {}, p .limit )
254- p .stats = & safeStats {
255- v : stats.Stats {Limit : p .limit },
256- onChange : p .trace .OnChange ,
224+ }
225+
226+ // createItemWithContext handles the creation of an item with context handling
227+ func createItemWithContext [PT Item [T ], T any ](
228+ ctx context.Context ,
229+ p * Pool [PT , T ],
230+ createItem func (ctx context.Context ) (PT , error ),
231+ ch chan PT ,
232+ ) error {
233+ var (
234+ createCtx = xcontext .ValueOnly (ctx )
235+ cancelCreate context.CancelFunc
236+ )
237+
238+ if d := p .createTimeout ; d > 0 {
239+ createCtx , cancelCreate = xcontext .WithTimeout (createCtx , d )
240+ } else {
241+ createCtx , cancelCreate = xcontext .WithCancel (createCtx )
257242 }
243+ defer cancelCreate ()
258244
259- return p
245+ newItem , err := createItem (createCtx )
246+ if err != nil {
247+ return xerrors .WithStackTrace (err )
248+ }
249+
250+ needCloseItem := true
251+ defer func () {
252+ if needCloseItem {
253+ _ = p .closeItem (ctx , newItem )
254+ }
255+ }()
256+
257+ select {
258+ case <- p .done :
259+ return xerrors .WithStackTrace (errClosedPool )
260+ case <- ctx .Done ():
261+ p .mu .Lock ()
262+ defer p .mu .Unlock ()
263+
264+ if len (p .index ) < p .limit {
265+ p .idle = append (p .idle , newItem )
266+ p .index [newItem ] = struct {}{}
267+ p .stats .Index ().Inc ()
268+ needCloseItem = false
269+ }
270+
271+ return xerrors .WithStackTrace (ctx .Err ())
272+ case ch <- newItem :
273+ needCloseItem = false
274+
275+ return nil
276+ }
260277}
261278
262279func (p * Pool [PT , T ]) Stats () stats.Stats {
0 commit comments