From 5363981c2dc4368ab11a377b0c4ce83b257e4301 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Tue, 5 Sep 2023 03:21:10 +0800
Subject: [PATCH] add

---
 .github/workflows/test.yml |  1 +
 ctx/Ctx.go                 | 77 ++++++++++++++++++++++++++++++++++++++
 ctx/Ctx_test.go            | 73 ++++++++++++++++++++++++++++++++++++
 3 files changed, 151 insertions(+)
 create mode 100644 ctx/Ctx.go
 create mode 100644 ctx/Ctx_test.go

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 71383dd..14eb277 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -39,6 +39,7 @@ jobs:
         go test -count 1 -timeout 10s -v -race github.com/qydysky/part/sql
         go test -count 1 -timeout 10s -v -race github.com/qydysky/part/rpc
         go test -count 1 -timeout 5s -v -race github.com/qydysky/part/component
+        go test -count 1 -timeout 15s -v -race github.com/qydysky/part/ctx
 
     - name: Set Release Name
       run: |
diff --git a/ctx/Ctx.go b/ctx/Ctx.go
new file mode 100644
index 0000000..a5da6e4
--- /dev/null
+++ b/ctx/Ctx.go
@@ -0,0 +1,77 @@
+package ctx
+
+import (
+	"context"
+	"errors"
+	"runtime"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	ErrWaitTo = errors.New("ErrWaitTo")
+	ErrToLess = errors.New("ErrToLess")
+)
+
+type Ctx struct {
+	Ctx context.Context
+	i32 *atomic.Int32
+	to  time.Duration
+}
+
+// ctx,done := WithWaitTo(..)
+//
+//	go func(){
+//			done1 := ctx.Wait()
+//			defer done1()
+//	}()
+//
+// done()// wait done1
+func WithWaitTo(sctx context.Context, to time.Duration) (ctx *Ctx, done func() error) {
+	if ctx, ok := sctx.(*Ctx); ok {
+		if ctx.to < to {
+			panic(ErrToLess)
+		}
+		ctx.i32.Add(1)
+	}
+	ctx = &Ctx{Ctx: sctx, i32: &atomic.Int32{}, to: to}
+	done = func() error {
+		<-ctx.Ctx.Done()
+		if ctx, ok := sctx.(*Ctx); ok {
+			defer ctx.i32.Add(-1)
+		}
+		be := time.Now()
+		for !ctx.i32.CompareAndSwap(0, -1) {
+			if time.Since(be) > to {
+				return ErrWaitTo
+			}
+			runtime.Gosched()
+		}
+		return nil
+	}
+	return
+}
+
+func (t Ctx) Deadline() (deadline time.Time, ok bool) {
+	return t.Ctx.Deadline()
+}
+
+func (t Ctx) Done() <-chan struct{} {
+	return t.Ctx.Done()
+}
+
+func (t Ctx) Err() error {
+	return t.Ctx.Err()
+}
+
+func (t Ctx) Value(key any) any {
+	return t.Ctx.Value(key)
+}
+
+func (t Ctx) Wait() (done func()) {
+	t.i32.Add(1)
+	<-t.Ctx.Done()
+	return func() {
+		t.i32.Add(-1)
+	}
+}
diff --git a/ctx/Ctx_test.go b/ctx/Ctx_test.go
new file mode 100644
index 0000000..3262877
--- /dev/null
+++ b/ctx/Ctx_test.go
@@ -0,0 +1,73 @@
+package ctx
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestMain(t *testing.T) {
+	ctx, _ := context.WithTimeout(context.Background(), time.Second)
+	ctx1, done := WithWaitTo(ctx, time.Second)
+	go func() {
+		done := ctx1.Wait()
+		defer done()
+	}()
+	if done() != nil {
+		t.Fatal()
+	}
+}
+
+func TestMain2(t *testing.T) {
+	ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
+	ctx1, done := WithWaitTo(ctx, time.Second)
+	go func() {
+		done := ctx1.Wait()
+		time.Sleep(time.Second * 2)
+		defer done()
+	}()
+	time.Sleep(time.Second)
+	if e := done(); !errors.Is(e, ErrWaitTo) {
+		t.Fatal(e)
+	}
+}
+
+func TestMain3(t *testing.T) {
+	ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
+	ctx1, done := WithWaitTo(ctx, time.Second)
+	go func() {
+		ctx2, done := WithWaitTo(ctx1, time.Second)
+		go func() {
+			done := ctx2.Wait()
+			defer done()
+		}()
+		if done() != nil {
+			t.Fail()
+		}
+	}()
+	time.Sleep(time.Second)
+	if done() != nil {
+		t.Fatal()
+	}
+}
+
+func TestMain4(t *testing.T) {
+	ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
+	ctx1, done := WithWaitTo(ctx, time.Second)
+	go func() {
+		ctx2, done := WithWaitTo(ctx1, time.Second)
+		go func() {
+			done := ctx2.Wait()
+			time.Sleep(time.Second * 2)
+			defer done()
+		}()
+		if e := done(); !errors.Is(e, ErrWaitTo) {
+			t.Fail()
+		}
+	}()
+	time.Sleep(time.Second)
+	if e := done(); !errors.Is(e, ErrWaitTo) {
+		t.Fail()
+	}
+}
-- 
2.39.2