From b1da1c50e553de3233a041e860714db02ebe8a17 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Mon, 27 Jan 2025 13:36:55 +0800
Subject: [PATCH] 1 (#18)

---
 ctx/Ctx.go      |  7 +++++++
 ctx/Ctx_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/ctx/Ctx.go b/ctx/Ctx.go
index 62c056c..361c44c 100644
--- a/ctx/Ctx.go
+++ b/ctx/Ctx.go
@@ -13,6 +13,7 @@ var (
 	ptr            = &struct{}{}
 	ErrWaitTo      = errors.New("ErrWaitTo")
 	ErrNothingWait = errors.New("ErrNothingWait")
+	ErrDoneCalled  = errors.New("ErrDoneCalled")
 )
 
 type Ctx struct {
@@ -44,7 +45,13 @@ func WithWait(sctx context.Context, planNum int32, to ...time.Duration) (dctx co
 
 	var doneWait context.CancelFunc
 	dctx, doneWait = context.WithCancel(ctx.Ctx)
+
+	var oncef atomic.Bool
 	done = func() error {
+		if !oncef.CompareAndSwap(false, true) {
+			return ErrDoneCalled
+		}
+
 		doneWait()
 		if ctxp, ok := sctx.Value(ptr).(*Ctx); ok {
 			defer func() {
diff --git a/ctx/Ctx_test.go b/ctx/Ctx_test.go
index f75a775..f15c07f 100644
--- a/ctx/Ctx_test.go
+++ b/ctx/Ctx_test.go
@@ -30,6 +30,31 @@ func TestMain(t *testing.T) {
 }
 
 func TestMain5(t *testing.T) {
+	ctx1, done := WithWait(context.Background(), 1)
+	t0 := time.Now()
+	go func() {
+		ctx2, done1 := WaitCtx(ctx1)
+		defer done1()
+		<-ctx2.Done()
+		if time.Since(t0) < time.Millisecond*100 {
+			t.Fail()
+		}
+		time.Sleep(time.Second)
+	}()
+	time.Sleep(time.Second)
+	t1 := time.Now()
+	if done() != nil {
+		t.Fatal()
+	}
+	if time.Since(t1) < time.Second {
+		t.Fail()
+	}
+	if !errors.Is(done(), ErrDoneCalled) {
+		t.Fatal()
+	}
+}
+
+func TestMain7(t *testing.T) {
 	ctx1, done := WithWait(context.Background(), 1, time.Second*2)
 	go func() {
 		ctx2, done1 := WaitCtx(ctx1)
-- 
2.39.2