update #92 benckmark

bench.go

@@ -8,6 +8,7 @@ type Bench struct {
 	durations  time.Duration
 	number     int
 	concurrent int
+	rate       int
 
 	g *routerGroup
 }
@@ -22,6 +23,11 @@ func (b *Bench) Number(n int) *Bench {
 	return b
 }
 
+func (b *Bench) Rate(rate int) *Bench {
+	b.rate = rate
+	return b
+}
+
 func (b *Bench) Durations(d time.Duration) *Bench {
 	b.durations = d
 	return b

bench/report.go

@@ -1,13 +1,14 @@
 package bench
 
 import (
+	"sync/atomic"
 	"time"
 )
 
 // 数据字段，每个字段都用于显示
 type report struct {
 	Concurrency   int
-	Failed        int
+	Failed        int32
 	Tps           float64
 	Total         time.Duration
 	Kbs           float64
@@ -27,5 +28,5 @@ type Report struct {
 }
 
 func (r *Report) AddFail() {
-	atomic.AddInt32(&r.Failed)
+	atomic.AddInt32(&r.Failed, 1)
 }

bench/task.go

@@ -0,0 +1,115 @@
+package bench
+
+import (
+	"os"
+	"os/signal"
+	"sync"
+	"time"
+)
+
+type Task struct {
+	Duration   time.Duration //压测时间
+	Number     int           //压测次数
+	Concurrent int           //并发数
+	Rate       int           //压测频率
+
+	work chan struct{}
+
+	ok bool
+
+	wg sync.WaitGroup
+}
+
+func (t *Task) Init() {
+	t.work = make(chan struct{})
+	t.ok = true
+}
+
+func (t *Task) Producer() {
+	if t.ok == false {
+		panic("task must be init")
+	}
+
+	work := t.work
+	// 控制压测时间
+	if t.Duration > 0 {
+		tk := time.NewTicker(t.Duration)
+		go func() {
+			defer close(work)
+			for {
+				select {
+				case <-tk.C:
+				case work <- struct{}{}:
+				}
+			}
+		}()
+
+	}
+
+	go func() {
+		defer close(work)
+
+		switch {
+		case t.Number == 0:
+			return
+		case t.Number > 0:
+			for i, n := 0, t.Number; i < n; i++ {
+				work <- struct{}{}
+			}
+		default: // t.Number < 0
+			for {
+				work <- struct{}{}
+			}
+		}
+
+	}()
+
+}
+
+func (t *Task) RunMain() {
+	sig := make(chan os.Signal, 1)
+	signal.Notify(sig, os.Interrupt)
+
+	interval := 0
+	work := t.work
+	wg := &t.wg
+
+	if t.Rate > 0 {
+		interval = int(time.Second) / t.Rate
+	}
+
+	begin := time.Now()
+	if interval > 0 {
+		oldwork := work
+		count := 0
+		work = make(chan struct{}, 1)
+
+		wg.Add(1)
+		go func() {
+			defer func() {
+				close(work)
+				wg.Done()
+			}()
+
+			for {
+				next := begin.Add(time.Duration(count * interval))
+				time.Sleep(next.Sub(time.Now()))
+				select {
+				case _, ok := <-oldwork:
+					if !ok {
+						return
+					}
+				default:
+				}
+
+				work <- struct{}{}
+				count++
+			}
+		}()
+
+	}
+
+	for i, c := 0, t.Concurrent; i < c; i++ {
+		//wg.Add(1)
+	}
+}