update

Author: admpub

goroutine.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"os/signal"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -43,14 +44,19 @@ func Notify(sig ...os.Signal) chan os.Signal {
 	return terminate
 }
 
-func NewDelayOnce(delay time.Duration, timeout time.Duration) *DelayOnce {
+func NewDelayOnce(delay time.Duration, timeout time.Duration, debugMode ...bool) *DelayOnce {
 	if timeout <= delay {
 		panic(`timeout must be greater than delay`)
 	}
+	var debug bool
+	if len(debugMode) > 0 {
+		debug = debugMode[0]
+	}
 	return &DelayOnce{
 		mp:      sync.Map{},
 		delay:   delay,
 		timeout: timeout,
+		debug:   debug,
 	}
 }
 
@@ -65,12 +71,14 @@ type DelayOnce struct {
 	mp      sync.Map
 	delay   time.Duration
 	timeout time.Duration
+	debug   bool
 }
 
 type eventSession struct {
 	cancel context.CancelFunc
 	time   time.Time
 	mutex  sync.RWMutex
+	stop   chan struct{}
 }
 
 func (e *eventSession) Renew(t time.Time) {
@@ -86,58 +94,109 @@ func (e *eventSession) Time() time.Time {
 	return t
 }
 
-func (d *DelayOnce) checkAndStore(parentCtx context.Context, key string) (exit bool, ctx context.Context) {
+func (e *eventSession) Cancel() <-chan struct{} {
+	e.cancel()
+	return e.stop
+}
+
+func (d *DelayOnce) checkAndStore(parentCtx context.Context, key string) (*eventSession, context.Context) {
 	v, loaded := d.mp.Load(key)
 	if loaded {
 		session := v.(*eventSession)
 		if time.Since(session.Time()) < d.timeout { // 超过 d.timeout 后重新处理，d.timeout 内记录当前时间
 			session.Renew(time.Now())
 			d.mp.Store(key, session)
-			return true, nil
+			return nil, nil
+		}
+
+		if d.debug {
+			log.Println(`[DelayOnce] cancel -------------> ` + key)
+		}
+
+		<-session.Cancel()
+
+		if d.debug {
+			log.Println(`[DelayOnce] canceled -------------> ` + key)
 		}
-		session.cancel()
 	}
-	var cancel context.CancelFunc
-	ctx, cancel = context.WithCancel(parentCtx)
-	d.mp.Store(key, &eventSession{
+	ctx, cancel := context.WithCancel(parentCtx)
+	session := &eventSession{
 		cancel: cancel,
 		time:   time.Now(),
-	})
-	return false, ctx
+		stop:   make(chan struct{}, 1),
+	}
+	d.mp.Store(key, session)
+	return session, ctx
 }
 
 func (d *DelayOnce) Do(parentCtx context.Context, key string, f func() error) (isNew bool) {
-	exit, ctx := d.checkAndStore(parentCtx, key)
-	if exit {
+	session, ctx := d.checkAndStore(parentCtx, key)
+	if session == nil {
 		return false
 	}
 	go func(key string) {
 		for {
 			t := time.NewTicker(time.Second)
 			defer t.Stop()
 			select {
-			case <-ctx.Done():
+			case <-ctx.Done(): // 如果先进入“<-t.C”分支，会等“<-t.C”分支内的代码执行完毕后才有机会执行本分支
 				d.mp.Delete(key)
+				session.stop <- struct{}{}
+				close(session.stop)
+				if d.debug {
+					log.Println(`[DelayOnce] close -------------> ` + key)
+				}
 				return
 			case <-t.C:
-				if err := d.exec(key, f); err != nil {
-					log.Println(key+`:`, err)
+				if time.Since(session.Time()) > d.delay { // 时间超过d.delay才触发
+					err := f()
+					session.Cancel()
+					if err != nil {
+						log.Println(key+`:`, err)
+					}
 				}
 			}
 		}
 	}(key)
 	return true
 }
 
-func (d *DelayOnce) exec(key string, f func() error) (err error) {
-	v, ok := d.mp.Load(key)
-	if !ok {
-		return
-	}
-	session := v.(*eventSession)
-	if time.Since(session.Time()) > d.delay { // 时间超过d.delay才触发
-		err = f()
+func (d *DelayOnce) DoWithState(parentCtx context.Context, key string, f func(func() bool) error) (isNew bool) {
+	session, ctx := d.checkAndStore(parentCtx, key)
+	if session == nil {
+		return false
 	}
-	session.cancel()
-	return
+	go func(key string) {
+		var state int32
+		isAbort := func() bool {
+			return atomic.LoadInt32(&state) > 0
+		}
+		go func() {
+			<-ctx.Done()
+			atomic.AddInt32(&state, 1)
+		}()
+		for {
+			t := time.NewTicker(time.Second)
+			defer t.Stop()
+			select {
+			case <-ctx.Done(): // 如果先进入“<-t.C”分支，会等“<-t.C”分支内的代码执行完毕后才有机会执行本分支
+				d.mp.Delete(key)
+				session.stop <- struct{}{}
+				close(session.stop)
+				if d.debug {
+					log.Println(`[DelayOnce] close -------------> ` + key)
+				}
+				return
+			case <-t.C:
+				if time.Since(session.Time()) > d.delay { // 时间超过d.delay才触发
+					err := f(isAbort)
+					session.Cancel()
+					if err != nil {
+						log.Println(key+`:`, err)
+					}
+				}
+			}
+		}
+	}(key)
+	return true
 }

goroutine_test.go

@@ -3,6 +3,7 @@ package com
 import (
 	"context"
 	"fmt"
+	"log"
 	"sync"
 	"testing"
 	"time"
@@ -18,15 +19,15 @@ func TestLoop(t *testing.T) {
 	cancel()
 }
 
-func TestDelayOnce(t *testing.T) {
-	d := NewDelayOnce(time.Second*2, time.Hour)
+func TestDelayOnceNormal(t *testing.T) {
+	d := NewDelayOnce(time.Second*2, time.Hour, true)
 	ctx := context.TODO()
 	wg := sync.WaitGroup{}
 	for i := 0; i < 10; i++ {
-		fmt.Println(`Trigger`, time.Now())
-		isNew := d.Do(ctx, `key`, func() error {
+		log.Println(`Trigger key_normal`)
+		isNew := d.Do(ctx, `key_normal`, func() error {
 			defer wg.Done()
-			fmt.Println(`Execute`, time.Now())
+			log.Println(`Execute key_normal`)
 			return nil
 		})
 		if isNew {
@@ -36,3 +37,30 @@ func TestDelayOnce(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestDelayOnceTimeout(t *testing.T) {
+	d := NewDelayOnce(time.Second*2, time.Second*5, true)
+	ctx := context.TODO()
+	wg := sync.WaitGroup{}
+	for i := 0; i < 10; i++ {
+		log.Println(`Trigger key_timeout`)
+		isNew := d.DoWithState(ctx, `key_timeout`, func(isAbort func() bool) error {
+			defer wg.Done()
+			for i := 0; i < 4; i++ {
+				if isAbort() {
+					log.Println(`------> Stop key_timeout`)
+					return nil
+				}
+				log.Println(`Execute key_timeout`, i)
+				time.Sleep(time.Second * 5)
+			}
+			log.Println(`Execute key_timeout`)
+			return nil
+		})
+		if isNew {
+			wg.Add(1)
+		}
+		time.Sleep(time.Second * 6)
+	}
+	wg.Wait()
+}