修复context问题和添加变量支持

cmd/main.go

@@ -7,6 +7,7 @@ import (
 	"git.nspix.com/golang/kos/entry/cli"
 	"git.nspix.com/golang/kos/entry/http"
 	httpkg "net/http"
+	"time"
 
 	"git.nspix.com/golang/kos"
 )
@@ -28,18 +29,28 @@ func (s *subServer) Start(ctx context.Context) (err error) {
 	kos.Http().Handle(httpkg.MethodGet, "/hello", func(ctx *http.Context) (err error) {
 		return ctx.Success("Hello World")
 	})
+
 	kos.Command().Handle("/test", "test command", func(ctx *cli.Context) (err error) {
 		return ctx.Success([][]string{
 			[]string{"NAME", "AGE"},
 			[]string{"SSS", "aaa"},
 		})
 	})
+	
 	kos.Command().Handle("/users", "test command", func(ctx *cli.Context) (err error) {
 		return ctx.Success([]*users{
 			{Name: "Zhan", Age: 10, Tags: []string{"a", "b"}},
 			{Name: "Lisi", Age: 15, Tags: []string{"c", "d"}},
 		})
 	})
+
+	kos.Command().Handle("/ctx", "context test", func(ctx *cli.Context) (err error) {
+		select {
+		case <-ctx.Context().Done():
+		case <-time.After(time.Second * 2):
+		}
+		return ctx.Success("OK")
+	})
 	return
 }
 

entry/cli/context.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"sync"
 )
 
 type Context struct {
@@ -13,6 +14,8 @@ type Context struct {
 	ctx       context.Context
 	wc        io.WriteCloser
 	params    map[string]string
+	locker    sync.RWMutex
+	variables map[string]any
 	args      []string
 }
 
@@ -23,6 +26,7 @@ func (ctx *Context) reset(id int64, wc io.WriteCloser) {
 	ctx.ctx = context.Background()
 	ctx.args = make([]string, 0)
 	ctx.params = make(map[string]string)
+	ctx.variables = make(map[string]any)
 }
 
 func (ctx *Context) setArgs(args []string) {
@@ -42,7 +46,7 @@ func (ctx *Context) setContext(c context.Context) {
 }
 
 func (ctx *Context) Context() context.Context {
-	return ctx.Context()
+	return ctx.ctx
 }
 
 func (ctx *Context) Argument(index int) string {
@@ -59,6 +63,22 @@ func (ctx *Context) Param(s string) string {
 	return ""
 }
 
+func (ctx *Context) SetValue(name string, value any) {
+	ctx.locker.Lock()
+	if ctx.variables == nil {
+		ctx.variables = make(map[string]any)
+	}
+	ctx.variables[name] = value
+	ctx.locker.Unlock()
+}
+
+func (ctx *Context) GetValue(name string) (val any, ok bool) {
+	ctx.locker.RLock()
+	defer ctx.locker.RUnlock()
+	val, ok = ctx.variables[name]
+	return
+}
+
 func (ctx *Context) Success(v any) (err error) {
 	return ctx.send(responsePayload{Type: PacketTypeCommand, Data: v})
 }

entry/cli/server.go

@@ -44,7 +44,7 @@ func (svr *Server) releaseContext(ctx *Context) {
 	ctxPool.Put(ctx)
 }
 
-func (svr *Server) handle(ctx *Context, frame *Frame) (err error) {
+func (svr *Server) execute(ctx *Context, frame *Frame) (err error) {
 	var (
 		params map[string]string
 		tokens []string
@@ -139,7 +139,7 @@ func (svr *Server) process(conn net.Conn) {
 				break
 			}
 		case PacketTypeCommand:
-			if err = svr.handle(ctx, frame); err != nil {
+			if err = svr.execute(ctx, frame); err != nil {
 				break
 			}
 		default: