优化代码

Author: FishGoddess

FUTURE.md

@@ -1,5 +1,10 @@
 ## ✈️ 未来版本的新特性 (Features in future versions)
 
+### v0.2.x
+
+* [x] 优化代码
+* [ ] 增加生命周期检测，资源检查回调
+
 ### v0.1.x
 
 * [x] 实现基础的控制池功能

HISTORY.md

@@ -1,5 +1,11 @@
 ## 📜 历史版本的特性介绍 (Features in old versions)
 
+### v0.2.0
+
+> 此版本发布于 2025-04-12
+
+* 优化代码
+
 ### v0.1.1
 
 > 此版本发布于 2024-03-07

_icons/coverage.svg

@@ -4,20 +4,36 @@
 
 package rego
 
-import "context"
+import (
+	"context"
+	"errors"
+)
+
+var (
+	ErrPoolIsFull   = errors.New("rego: pool is full")
+	ErrPoolIsClosed = errors.New("rego: pool is closed")
+)
 
 type config struct {
 	fastFailed bool
 
-	newPoolFullErrFunc   func(ctx context.Context) error
-	newPoolClosedErrFunc func(ctx context.Context) error
+	newPoolFullErr   func(ctx context.Context) error
+	newPoolClosedErr func(ctx context.Context) error
 }
 
 func newDefaultConfig() *config {
+	newPoolFullErr := func(_ context.Context) error {
+		return ErrPoolIsFull
+	}
+
+	newPoolClosedErr := func(_ context.Context) error {
+		return ErrPoolIsClosed
+	}
+
 	conf := &config{
-		fastFailed:           false,
-		newPoolFullErrFunc:   nil,
-		newPoolClosedErrFunc: nil,
+		fastFailed:       false,
+		newPoolFullErr:   newPoolFullErr,
+		newPoolClosedErr: newPoolClosedErr,
 	}
 
 	return conf
@@ -39,13 +55,17 @@ func WithFastFailed() Option {
 // WithPoolFullErr sets newPoolFullErr to config.
 func WithPoolFullErr(newPoolFullErr func(ctx context.Context) error) Option {
 	return func(conf *config) {
-		conf.newPoolFullErrFunc = newPoolFullErr
+		if newPoolFullErr != nil {
+			conf.newPoolFullErr = newPoolFullErr
+		}
 	}
 }
 
 // WithPoolClosedErr sets newPoolClosedErr to config.
 func WithPoolClosedErr(newPoolClosedErr func(ctx context.Context) error) Option {
 	return func(conf *config) {
-		conf.newPoolClosedErrFunc = newPoolClosedErr
+		if newPoolClosedErr != nil {
+			conf.newPoolClosedErr = newPoolClosedErr
+		}
 	}
 }

option.go

@@ -10,7 +10,7 @@ import (
 	"testing"
 )
 
-// go test -v -cover -count=1 -test.cpu=1 -run=^TestWithFastFailed$
+// go test -v -cover -run=^TestWithFastFailed$
 func TestWithFastFailed(t *testing.T) {
 	conf := &config{fastFailed: false}
 	WithFastFailed()(conf)
@@ -20,30 +20,42 @@ func TestWithFastFailed(t *testing.T) {
 	}
 }
 
-// go test -v -cover -count=1 -test.cpu=1 -run=^TestWithPoolFullErr$
+// go test -v -cover -run=^TestWithPoolFullErr$
 func TestWithPoolFullErr(t *testing.T) {
 	newPoolFullErr := func(ctx context.Context) error {
 		return nil
 	}
 
-	conf := &config{newPoolFullErrFunc: nil}
+	conf := &config{newPoolFullErr: nil}
 	WithPoolFullErr(newPoolFullErr)(conf)
 
-	if fmt.Sprintf("%p", conf.newPoolFullErrFunc) != fmt.Sprintf("%p", newPoolFullErr) {
-		t.Fatalf("conf.newPoolFullErrFunc %p is wrong", conf.newPoolFullErrFunc)
+	if fmt.Sprintf("%p", conf.newPoolFullErr) != fmt.Sprintf("%p", newPoolFullErr) {
+		t.Fatalf("conf.newPoolFullErr %p is wrong", conf.newPoolFullErr)
+	}
+
+	WithPoolFullErr(nil)(conf)
+
+	if fmt.Sprintf("%p", conf.newPoolFullErr) != fmt.Sprintf("%p", newPoolFullErr) {
+		t.Fatalf("conf.newPoolFullErr %p is wrong", conf.newPoolFullErr)
 	}
 }
 
-// go test -v -cover -count=1 -test.cpu=1 -run=^TestWithPoolClosedErr$
+// go test -v -cover -run=^TestWithPoolClosedErr$
 func TestWithPoolClosedErr(t *testing.T) {
 	newPoolClosedErr := func(ctx context.Context) error {
 		return nil
 	}
 
-	conf := &config{newPoolClosedErrFunc: nil}
+	conf := &config{newPoolClosedErr: nil}
 	WithPoolClosedErr(newPoolClosedErr)(conf)
 
-	if fmt.Sprintf("%p", conf.newPoolClosedErrFunc) != fmt.Sprintf("%p", newPoolClosedErr) {
-		t.Fatalf("conf.newPoolClosedErrFunc %p is wrong", conf.newPoolClosedErrFunc)
+	if fmt.Sprintf("%p", conf.newPoolClosedErr) != fmt.Sprintf("%p", newPoolClosedErr) {
+		t.Fatalf("conf.newPoolClosedErr %p is wrong", conf.newPoolClosedErr)
+	}
+
+	WithPoolClosedErr(nil)(conf)
+
+	if fmt.Sprintf("%p", conf.newPoolClosedErr) != fmt.Sprintf("%p", newPoolClosedErr) {
+		t.Fatalf("conf.newPoolClosedErr %p is wrong", conf.newPoolClosedErr)
 	}
 }

option_test.go

@@ -6,15 +6,9 @@ package rego
 
 import (
 	"context"
-	"errors"
 	"sync"
 )
 
-var (
-	ErrPoolIsFull   = errors.New("rego: pool is full")
-	ErrPoolIsClosed = errors.New("rego: pool is closed")
-)
-
 type Status struct {
 	// Limit is the limit of acquired resources.
 	Limit uint64 `json:"limit"`
@@ -57,6 +51,10 @@ type Pool[T any] struct {
 }
 
 func New[T any](limit uint64, acquire AcquireFunc[T], release ReleaseFunc[T], opts ...Option) *Pool[T] {
+	if limit <= 0 {
+		panic("rego: limit can't be less than 0")
+	}
+
 	if acquire == nil || release == nil {
 		panic("rego: acquire or release func can't be nil")
 	}
@@ -94,22 +92,6 @@ func (p *Pool[T]) Put(resource T) error {
 	}
 }
 
-func (p *Pool[T]) newPoolFullErr(ctx context.Context) error {
-	if p.conf.newPoolFullErrFunc == nil {
-		return ErrPoolIsFull
-	}
-
-	return p.conf.newPoolFullErrFunc(ctx)
-}
-
-func (p *Pool[T]) newPoolClosedErr(ctx context.Context) error {
-	if p.conf.newPoolClosedErrFunc == nil {
-		return ErrPoolIsClosed
-	}
-
-	return p.conf.newPoolClosedErrFunc(ctx)
-}
-
 func (p *Pool[T]) tryToTake() (resource T, ok bool) {
 	select {
 	case resource = <-p.resources:
@@ -144,7 +126,7 @@ func (p *Pool[T]) Take(ctx context.Context) (resource T, err error) {
 	if p.closed {
 		p.lock.Unlock()
 
-		return resource, p.newPoolClosedErr(ctx)
+		return resource, p.conf.newPoolClosedErr(ctx)
 	}
 
 	var ok bool
@@ -174,7 +156,7 @@ func (p *Pool[T]) Take(ctx context.Context) (resource T, err error) {
 	if p.conf.fastFailed {
 		p.lock.Unlock()
 
-		return resource, p.newPoolFullErr(ctx)
+		return resource, p.conf.newPoolFullErr(ctx)
 	}
 
 	p.waiting++
@@ -205,19 +187,16 @@ func (p *Pool[T]) Status() Status {
 }
 
 func (p *Pool[T]) releaseResources() error {
-	for i := uint64(0); i < p.acquired; i++ {
-		resource := <-p.resources
-		if err := p.release(resource); err != nil {
-			return err
+	for {
+		select {
+		case resource := <-p.resources:
+			if err := p.release(resource); err != nil {
+				return err
+			}
+		default:
+			return nil
 		}
 	}
-	// for resource := range p.resources {
-	// 	if err := p.release(resource); err != nil {
-	// 		return err
-	// 	}
-	// }
-
-	return nil
 }
 
 // Close closes pool and releases all resources.

pool.go

@@ -12,7 +12,7 @@ import (
 	"time"
 )
 
-// go test -v -cover -count=1 -test.cpu=1 -run=^TestWithFastFailed$
+// go test -v -cover -run=^TestWithFastFailed$
 func TestPool(t *testing.T) {
 	limit := int64(16)
 	acquireLimit := int64(0)