mirror of
https://github.com/criyle/go-judge.git
synced 2025-11-04 14:50:02 +08:00
stream: add websocket transport layer support
This commit is contained in:
parent
a52f1360cd
commit
aa41950f89
23
README.cn.md
23
README.cn.md
@ -528,6 +528,29 @@ reboot
|
|||||||
- `-force-gc-target` 默认 `20m`, 堆内存使用超过该值是强制垃圾收集和归还内存
|
- `-force-gc-target` 默认 `20m`, 堆内存使用超过该值是强制垃圾收集和归还内存
|
||||||
- `-force-gc-interval` 默认 `5s`, 为后台线程检查的频繁程度
|
- `-force-gc-interval` 默认 `5s`, 为后台线程检查的频繁程度
|
||||||
|
|
||||||
|
### WebSocket 流接口
|
||||||
|
|
||||||
|
WebSocket 流接口是用于运行一个程序,同时和它的输入输出进行交互。所有的消息都应该使用 WebSocket 的 binary 格式来发送来避免兼容性问题。
|
||||||
|
|
||||||
|
```text
|
||||||
|
+--------+--------+---...
|
||||||
|
| 类型 | 载荷 ...
|
||||||
|
+--------|--------+---...
|
||||||
|
请求:
|
||||||
|
请求类型 =
|
||||||
|
1 - 运行请求 (载荷 = JSON 编码的请求体)
|
||||||
|
2 - 设置终端窗口大小 (载荷 = JSON 编码的请求体)
|
||||||
|
3 - 输入 (载荷 = 1 字节 (4 位的 命令下标 + 4 位的 文件描述符) + 输入内容)
|
||||||
|
4 - 取消 (没有载荷)
|
||||||
|
|
||||||
|
响应:
|
||||||
|
响应类型 =
|
||||||
|
1 - 运行结果 (载荷 = JSON 编码的运行结果)
|
||||||
|
2 - 输出 (载荷 = 1 字节 (4 位的 命令下标 + 4 位的 文件描述符) + 输入内容)
|
||||||
|
```
|
||||||
|
|
||||||
|
任何的不完整,或者不合法的消息会被认为是错误,并终止运行。
|
||||||
|
|
||||||
### 压力测试
|
### 压力测试
|
||||||
|
|
||||||
使用 `wrk` 和 `t.lua`: `wrk -s t.lua -c 1 -t 1 -d 30s --latency http://localhost:5050/run`.
|
使用 `wrk` 和 `t.lua`: `wrk -s t.lua -c 1 -t 1 -d 30s --latency http://localhost:5050/run`.
|
||||||
|
|||||||
@ -21,13 +21,13 @@ type grpcWrapper struct {
|
|||||||
sc pb.Executor_ExecStreamClient
|
sc pb.Executor_ExecStreamClient
|
||||||
}
|
}
|
||||||
|
|
||||||
func newGrpc(args []string, srvAddr *string) Stream {
|
func newGrpc(args []string, srvAddr string) Stream {
|
||||||
token := os.Getenv("TOKEN")
|
token := os.Getenv("TOKEN")
|
||||||
opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
|
opts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
|
||||||
if token != "" {
|
if token != "" {
|
||||||
opts = append(opts, grpc.WithPerRPCCredentials(newTokenAuth(token)))
|
opts = append(opts, grpc.WithPerRPCCredentials(newTokenAuth(token)))
|
||||||
}
|
}
|
||||||
conn, err := grpc.Dial(*srvAddr, opts...)
|
conn, err := grpc.Dial(srvAddr, opts...)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Fatalln("client", err)
|
log.Fatalln("client", err)
|
||||||
}
|
}
|
||||||
|
|||||||
@ -16,7 +16,9 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
srvAddr = flag.String("srv-addr", "localhost:5051", "GRPC server addr")
|
transport = flag.String("transport", "websocket", "defines transport layer (websocket / grpc)")
|
||||||
|
wsURL = flag.String("ws-url", "ws://localhost:5050/stream", "HTTP server url")
|
||||||
|
grpcAddr = flag.String("grpc-addr", "localhost:5051", "GRPC server addr")
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
@ -33,6 +35,7 @@ var env = []string{
|
|||||||
"TERM=" + os.Getenv("TERM"),
|
"TERM=" + os.Getenv("TERM"),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Stream defines the transport layer for stream execution
|
||||||
type Stream interface {
|
type Stream interface {
|
||||||
Send(*stream.Request) error
|
Send(*stream.Request) error
|
||||||
Recv() (*stream.Response, error)
|
Recv() (*stream.Response, error)
|
||||||
@ -44,8 +47,16 @@ func main() {
|
|||||||
if len(args) == 0 {
|
if len(args) == 0 {
|
||||||
args = []string{"/bin/bash"}
|
args = []string{"/bin/bash"}
|
||||||
}
|
}
|
||||||
w := newGrpc(args, srvAddr)
|
var s Stream
|
||||||
r, err := run(w, args)
|
switch *transport {
|
||||||
|
case "websocket":
|
||||||
|
s = newWebsocket(args, *wsURL)
|
||||||
|
case "grpc":
|
||||||
|
s = newGrpc(args, *grpcAddr)
|
||||||
|
default:
|
||||||
|
log.Fatalln("invalid transport: ", *transport)
|
||||||
|
}
|
||||||
|
r, err := run(s, args)
|
||||||
log.Printf("finished: %+v %v", r, err)
|
log.Printf("finished: %+v %v", r, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
104
cmd/go-judge-shell/websocket.go
Normal file
104
cmd/go-judge-shell/websocket.go
Normal file
@ -0,0 +1,104 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
|
||||||
|
"github.com/criyle/go-judge/cmd/go-judge/model"
|
||||||
|
"github.com/criyle/go-judge/cmd/go-judge/stream"
|
||||||
|
"github.com/gorilla/websocket"
|
||||||
|
)
|
||||||
|
|
||||||
|
var _ Stream = &websocketStream{}
|
||||||
|
|
||||||
|
type websocketStream struct {
|
||||||
|
conn *websocket.Conn
|
||||||
|
}
|
||||||
|
|
||||||
|
func newWebsocket(args []string, wsURL string) Stream {
|
||||||
|
header := make(http.Header)
|
||||||
|
token := os.Getenv("TOKEN")
|
||||||
|
if token != "" {
|
||||||
|
header.Add("Authorization", "Bearer "+token)
|
||||||
|
}
|
||||||
|
conn, _, err := websocket.DefaultDialer.Dial(wsURL, header)
|
||||||
|
if err != nil {
|
||||||
|
log.Fatalln("ws connect: ", err)
|
||||||
|
}
|
||||||
|
log.Println("start", args)
|
||||||
|
return &websocketStream{conn: conn}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Recv implements Stream.
|
||||||
|
func (s *websocketStream) Recv() (*stream.Response, error) {
|
||||||
|
_, r, err := s.conn.ReadMessage()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if len(r) == 0 {
|
||||||
|
return nil, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
resp := new(stream.Response)
|
||||||
|
switch r[0] {
|
||||||
|
case 1:
|
||||||
|
resp.Response = new(model.Response)
|
||||||
|
if err := json.Unmarshal(r[1:], resp.Response); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
case 2:
|
||||||
|
if len(r) < 2 {
|
||||||
|
return nil, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
resp.Output = new(stream.OutputResponse)
|
||||||
|
resp.Output.Index = int(r[1]>>4) & 0xf
|
||||||
|
resp.Output.Fd = int(r[1]) & 0xf
|
||||||
|
resp.Output.Content = r[2:]
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("invalid type code: %d", r[0])
|
||||||
|
}
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Send implements Stream.
|
||||||
|
func (s *websocketStream) Send(req *stream.Request) error {
|
||||||
|
w, err := s.conn.NextWriter(websocket.BinaryMessage)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer w.Close()
|
||||||
|
|
||||||
|
switch {
|
||||||
|
case req.Request != nil:
|
||||||
|
if _, err := w.Write([]byte{1}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := json.NewEncoder(w).Encode(req.Request); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
case req.Resize != nil:
|
||||||
|
if _, err := w.Write([]byte{2}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := json.NewEncoder(w).Encode(req.Resize); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
case req.Input != nil:
|
||||||
|
if _, err := w.Write([]byte{3, byte(req.Input.Index<<4 | req.Input.Fd)}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := w.Write(req.Input.Content); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
case req.Cancel != nil:
|
||||||
|
if _, err := w.Write([]byte{4}); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("invalid request")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@ -510,6 +510,7 @@ func generateHandleVersion(conf *config.Config, builderParam map[string]any) fun
|
|||||||
"pipeProxy": true,
|
"pipeProxy": true,
|
||||||
"symlink": true,
|
"symlink": true,
|
||||||
"addressSpaceLimit": true,
|
"addressSpaceLimit": true,
|
||||||
|
"stream": true,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -521,6 +522,7 @@ func generateHandleConfig(conf *config.Config, builderParam map[string]any) func
|
|||||||
"pipeProxy": true,
|
"pipeProxy": true,
|
||||||
"symlink": true,
|
"symlink": true,
|
||||||
"addressSpaceLimit": true,
|
"addressSpaceLimit": true,
|
||||||
|
"stream": true,
|
||||||
"fileStorePath": conf.Dir,
|
"fileStorePath": conf.Dir,
|
||||||
"runnerConfig": builderParam,
|
"runnerConfig": builderParam,
|
||||||
})
|
})
|
||||||
|
|||||||
@ -12,7 +12,10 @@ import (
|
|||||||
"github.com/criyle/go-judge/worker"
|
"github.com/criyle/go-judge/worker"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// FileError defines the location, file name and the detailed message for a failed file operation
|
||||||
type FileError = envexec.FileError
|
type FileError = envexec.FileError
|
||||||
|
|
||||||
|
// FileErrorType defines the location that file operation fails
|
||||||
type FileErrorType = envexec.FileErrorType
|
type FileErrorType = envexec.FileErrorType
|
||||||
|
|
||||||
// CmdFile defines file from multiple source including local / memory / cached or pipe collector
|
// CmdFile defines file from multiple source including local / memory / cached or pipe collector
|
||||||
|
|||||||
@ -40,12 +40,12 @@ type Response struct {
|
|||||||
|
|
||||||
// ResizeRequest defines resize operation to the virtual terminal
|
// ResizeRequest defines resize operation to the virtual terminal
|
||||||
type ResizeRequest struct {
|
type ResizeRequest struct {
|
||||||
Index int
|
Index int `json:"index,omitempty"`
|
||||||
Fd int
|
Fd int `json:"fd,omitempty"`
|
||||||
Rows int
|
Rows int `json:"rows,omitempty"`
|
||||||
Cols int
|
Cols int `json:"cols,omitempty"`
|
||||||
X int
|
X int `json:"x,omitempty"`
|
||||||
Y int
|
Y int `json:"y,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// InputRequest defines input operation from the remote
|
// InputRequest defines input operation from the remote
|
||||||
@ -97,15 +97,13 @@ func Start(baseCtx context.Context, s Stream, w worker.Worker, srcPrefix []strin
|
|||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
// stream in
|
// stream in
|
||||||
if len(streamIn) > 0 {
|
wg.Go(func() error {
|
||||||
wg.Go(func() error {
|
if err := streamInput(ctx, s, streamIn, execCancel); err != nil {
|
||||||
if err := streamInput(ctx, s, streamIn, execCancel); err != nil {
|
cancel()
|
||||||
cancel()
|
return err
|
||||||
return err
|
}
|
||||||
}
|
return nil
|
||||||
return nil
|
})
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
// stream out
|
// stream out
|
||||||
outCh := make(chan *OutputResponse, len(streamOut))
|
outCh := make(chan *OutputResponse, len(streamOut))
|
||||||
|
|||||||
120
cmd/go-judge/ws_executor/stream.go
Normal file
120
cmd/go-judge/ws_executor/stream.go
Normal file
@ -0,0 +1,120 @@
|
|||||||
|
package wsexecutor
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/criyle/go-judge/cmd/go-judge/model"
|
||||||
|
"github.com/criyle/go-judge/cmd/go-judge/stream"
|
||||||
|
"github.com/gorilla/websocket"
|
||||||
|
)
|
||||||
|
|
||||||
|
var _ stream.Stream = &streamWrapper{}
|
||||||
|
|
||||||
|
type streamWrapper struct {
|
||||||
|
ctx context.Context
|
||||||
|
conn *websocket.Conn
|
||||||
|
sendCh chan stream.Response
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *streamWrapper) sendLoop() {
|
||||||
|
conn := w.conn
|
||||||
|
defer conn.Close()
|
||||||
|
|
||||||
|
ticker := time.NewTicker(pingPeriod)
|
||||||
|
defer ticker.Stop()
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-w.ctx.Done():
|
||||||
|
return
|
||||||
|
case r := <-w.sendCh:
|
||||||
|
conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||||
|
switch {
|
||||||
|
case r.Response != nil:
|
||||||
|
w, err := conn.NextWriter(websocket.BinaryMessage)
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if _, err := w.Write([]byte{1}); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := json.NewEncoder(w).Encode(r.Response); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := w.Close(); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
conn.Close()
|
||||||
|
return
|
||||||
|
case r.Output != nil:
|
||||||
|
w, err := conn.NextWriter(websocket.BinaryMessage)
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if _, err := w.Write([]byte{2, byte(r.Output.Index<<4 | r.Output.Fd)}); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if _, err := w.Write(r.Output.Content); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := w.Close(); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
case <-ticker.C:
|
||||||
|
conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||||
|
if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *streamWrapper) Send(resp stream.Response) error {
|
||||||
|
select {
|
||||||
|
case <-w.ctx.Done():
|
||||||
|
return w.ctx.Err()
|
||||||
|
case w.sendCh <- resp:
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *streamWrapper) Recv() (*stream.Request, error) {
|
||||||
|
conn := w.conn
|
||||||
|
_, buf, err := conn.ReadMessage()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if len(buf) == 0 {
|
||||||
|
return nil, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
var req stream.Request
|
||||||
|
switch buf[0] {
|
||||||
|
case 1:
|
||||||
|
req.Request = new(model.Request)
|
||||||
|
if err := json.Unmarshal(buf[1:], req.Request); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
case 2:
|
||||||
|
req.Resize = new(stream.ResizeRequest)
|
||||||
|
if err := json.Unmarshal(buf[1:], req.Resize); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
case 3:
|
||||||
|
if len(buf) < 2 {
|
||||||
|
return nil, io.ErrUnexpectedEOF
|
||||||
|
}
|
||||||
|
req.Input = new(stream.InputRequest)
|
||||||
|
req.Input.Index = int(buf[1]>>4) & 0xf
|
||||||
|
req.Input.Fd = int(buf[1]) & 0xf
|
||||||
|
req.Input.Content = buf[2:]
|
||||||
|
case 4:
|
||||||
|
req.Cancel = new(struct{})
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("invalid type code: %d", buf[0])
|
||||||
|
}
|
||||||
|
return &req, nil
|
||||||
|
}
|
||||||
@ -8,6 +8,7 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/criyle/go-judge/cmd/go-judge/model"
|
"github.com/criyle/go-judge/cmd/go-judge/model"
|
||||||
|
"github.com/criyle/go-judge/cmd/go-judge/stream"
|
||||||
"github.com/criyle/go-judge/worker"
|
"github.com/criyle/go-judge/worker"
|
||||||
"github.com/gin-gonic/gin"
|
"github.com/gin-gonic/gin"
|
||||||
"github.com/gorilla/websocket"
|
"github.com/gorilla/websocket"
|
||||||
@ -55,6 +56,7 @@ type wsRequest struct {
|
|||||||
|
|
||||||
func (h *wsHandle) Register(r *gin.Engine) {
|
func (h *wsHandle) Register(r *gin.Engine) {
|
||||||
r.GET("/ws", h.handleWS)
|
r.GET("/ws", h.handleWS)
|
||||||
|
r.GET("/stream", h.handleStream)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *wsHandle) handleWS(c *gin.Context) {
|
func (h *wsHandle) handleWS(c *gin.Context) {
|
||||||
@ -178,6 +180,26 @@ func (h *wsHandle) handleWS(c *gin.Context) {
|
|||||||
}()
|
}()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (h *wsHandle) handleStream(c *gin.Context) {
|
||||||
|
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||||
|
if err != nil {
|
||||||
|
c.Error(err)
|
||||||
|
c.AbortWithStatusJSON(http.StatusBadRequest, err.Error())
|
||||||
|
return
|
||||||
|
}
|
||||||
|
conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||||
|
conn.SetPongHandler(func(string) error {
|
||||||
|
conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
ctx, cancel := context.WithCancel(context.TODO())
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
w := &streamWrapper{ctx: ctx, conn: conn, sendCh: make(chan stream.Response)}
|
||||||
|
go w.sendLoop()
|
||||||
|
stream.Start(ctx, w, h.worker, h.srcPrefix, h.logger)
|
||||||
|
}
|
||||||
|
|
||||||
type contextMap struct {
|
type contextMap struct {
|
||||||
m map[string]context.CancelFunc
|
m map[string]context.CancelFunc
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user