From a8b798c2770357552d9ba59b5351544eadcd793c Mon Sep 17 00:00:00 2001
From: storezhang <storezhang@gmail.com>
Date: Thu, 6 Apr 2023 17:47:18 +0800
Subject: [PATCH] =?UTF-8?q?feat(=E4=BB=BB=E5=8A=A1):=20=E5=A2=9E=E5=8A=A0?=
 =?UTF-8?q?=E5=85=A8=E5=B1=80=E9=99=90=E5=88=B6=E6=9D=A1=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 add_builder.go   |  8 +++++++-
 add_params.go    |  2 ++
 builder.go       |  4 ++++
 limit_builder.go | 27 ++++++++++++++++-----------
 limit_params.go  |  9 +++++++++
 limit_type.go    |  5 +++++
 params.go        |  1 +
 scheduler.go     |  8 ++++++++
 8 files changed, 52 insertions(+), 12 deletions(-)
 create mode 100644 limit_type.go

diff --git a/add_builder.go b/add_builder.go
index dc03a4f..4c09356 100644
--- a/add_builder.go
+++ b/add_builder.go
@@ -59,10 +59,16 @@ func (ab *addBuilder) Random() (builder *randomBuilder) {
 	return
 }
 
-func (ab *addBuilder) Limit() *limitBuilder {
+func (ab *addBuilder) Limit() *limitBuilder[*addBuilder] {
 	return newLimitBuilder(ab)
 }
 
+func (ab *addBuilder) Unlimited() *addBuilder {
+	ab.self.limit = newUnlimitedParams()
+
+	return ab
+}
+
 func (ab *addBuilder) Unique() *addBuilder {
 	ab.self.unique = true
 
diff --git a/add_params.go b/add_params.go
index 719d278..740c875 100644
--- a/add_params.go
+++ b/add_params.go
@@ -27,6 +27,8 @@ func newAddParams(scheduler *Scheduler, worker worker) (add *addParams) {
 func (ap *addParams) checkLimit(scheduler *Scheduler) (err error) {
 	if nil != ap.limit {
 		err = ap.limit.check(scheduler)
+	} else {
+		err = scheduler.checkLimit()
 	}
 
 	return
diff --git a/builder.go b/builder.go
index ab4158f..5e76e74 100644
--- a/builder.go
+++ b/builder.go
@@ -26,6 +26,10 @@ func (b *builder) Unique() *builder {
 	return b
 }
 
+func (b *builder) Limit() *limitBuilder[*builder] {
+	return newLimitBuilder(b)
+}
+
 func (b *builder) Build() *Scheduler {
 	return newScheduler(b.params)
 }
diff --git a/limit_builder.go b/limit_builder.go
index e69866b..4ea95aa 100644
--- a/limit_builder.go
+++ b/limit_builder.go
@@ -1,44 +1,49 @@
 package schedule
 
-type limitBuilder struct {
+type limitBuilder[T limitType] struct {
 	params  *limitParams
-	builder *addBuilder
+	builder T
 }
 
-func newLimitBuilder(builder *addBuilder) *limitBuilder {
-	return &limitBuilder{
+func newLimitBuilder[T limitType](builder T) *limitBuilder[T] {
+	return &limitBuilder[T]{
 		params:  newLimitParams(),
 		builder: builder,
 	}
 }
 
-func (lb *limitBuilder) Cpu(percent float64) *limitBuilder {
+func (lb *limitBuilder[T]) Cpu(percent float64) *limitBuilder[T] {
 	lb.params.cpu = percent
 
 	return lb
 }
 
-func (lb *limitBuilder) Memory(percent float64) *limitBuilder {
+func (lb *limitBuilder[T]) Memory(percent float64) *limitBuilder[T] {
 	lb.params.memory = percent
 
 	return lb
 }
 
-func (lb *limitBuilder) Process(count int) *limitBuilder {
+func (lb *limitBuilder[T]) Process(count int) *limitBuilder[T] {
 	lb.params.process = count
 
 	return lb
 }
 
-func (lb *limitBuilder) Max(max int) *limitBuilder {
+func (lb *limitBuilder[T]) Max(max int) *limitBuilder[T] {
 	lb.params.max = max
 
 	return lb
 }
 
-func (lb *limitBuilder) Build() (builder *addBuilder) {
-	lb.builder.self.limit = lb.params
-	builder = lb.builder
+func (lb *limitBuilder[T]) Build() (t T) {
+	switch target := any(lb.builder).(type) {
+	case *builder:
+		target.params.limit = lb.params
+	case *addBuilder:
+		target.self.limit = lb.params
+	}
+	t = lb.builder
 
 	return
 }
diff --git a/limit_params.go b/limit_params.go
index 5d88813..6e8fe55 100644
--- a/limit_params.go
+++ b/limit_params.go
@@ -27,6 +27,15 @@ func newLimitParams() *limitParams {
 	}
 }
 
+func newUnlimitedParams() *limitParams {
+	return &limitParams{
+		cpu:     math.MaxFloat64,
+		memory:  math.MaxFloat64,
+		process: math.MaxInt,
+		max:     math.MaxInt,
+	}
+}
+
 func (lp *limitParams) check(scheduler *Scheduler) (err error) {
 	if ce := lp.checkCount(scheduler); nil != ce {
 		err = ce
diff --git a/limit_type.go b/limit_type.go
new file mode 100644
index 0000000..60573df
--- /dev/null
+++ b/limit_type.go
@@ -0,0 +1,5 @@
+package schedule
+
+type limitType interface {
+	*builder | *addBuilder
+}
diff --git a/params.go b/params.go
index e5be861..bda94b0 100644
--- a/params.go
+++ b/params.go
@@ -7,6 +7,7 @@ import (
 type params struct {
 	logger simaqian.Logger
 	unique bool
+	limit  *limitParams
 }
 
 func newParams() *params {
diff --git a/scheduler.go b/scheduler.go
index 3b3ed8d..f3bb7ac 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -85,3 +85,11 @@ func (s *Scheduler) remove(id string) {
 		s.ids.Delete(id)
 	}
 }
+
+func (s *Scheduler) checkLimit() (err error) {
+	if nil != s.params.limit {
+		err = s.params.limit.check(s)
+	}
+
+	return
+}