refactor(envexec): abstract stream in and out
Some checks are pending
Build / Goreleaser (push) Waiting to run
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (amd64_v3, darwin) (push) Blocked by required conditions
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (amd64_v3, linux) (push) Blocked by required conditions
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (amd64_v3, windows) (push) Blocked by required conditions
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (arm64_v8.0, darwin) (push) Blocked by required conditions
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (arm64_v8.0, linux) (push) Blocked by required conditions
Build / Upload artifacts-${{ matrix.os }}-${{ matrix.arch }} (arm64_v8.0, windows) (push) Blocked by required conditions

This commit is contained in:
criyle 2025-05-29 19:14:05 +00:00
parent c5575ffe2b
commit 94f57a3553
6 changed files with 136 additions and 102 deletions

View File

@ -1,8 +1,8 @@
package stream package stream
import ( import (
"errors"
"fmt" "fmt"
"os"
"github.com/criyle/go-judge/envexec" "github.com/criyle/go-judge/envexec"
"github.com/criyle/go-judge/filestore" "github.com/criyle/go-judge/filestore"
@ -12,6 +12,8 @@ import (
var ( var (
_ worker.CmdFile = &fileStreamIn{} _ worker.CmdFile = &fileStreamIn{}
_ worker.CmdFile = &fileStreamOut{} _ worker.CmdFile = &fileStreamOut{}
errNoTTY = errors.New("stream is not a tty")
) )
type fileStreamIn struct { type fileStreamIn struct {
@ -30,15 +32,15 @@ func newFileStreamIn(index, fd int, hasTTY bool) *fileStreamIn {
} }
} }
func (f *fileStreamIn) GetTTY() *os.File { func (f *fileStreamIn) SetSize(s *envexec.TerminalSize) error {
if !f.hasTTY { if !f.hasTTY {
return nil return errNoTTY
} }
return f.stream.WritePipe() return f.stream.SetSize(s)
} }
func (f *fileStreamIn) Write(b []byte) (int, error) { func (f *fileStreamIn) Write(b []byte) (int, error) {
return f.stream.WritePipe().Write(b) return f.stream.Write(b)
} }
func (f *fileStreamIn) EnvFile(fs filestore.FileStore) (envexec.File, error) { func (f *fileStreamIn) EnvFile(fs filestore.FileStore) (envexec.File, error) {
@ -68,7 +70,7 @@ func newFileStreamOut(index, fd int) *fileStreamOut {
} }
func (f *fileStreamOut) Read(b []byte) (int, error) { func (f *fileStreamOut) Read(b []byte) (int, error) {
return f.stream.ReadPipe().Read(b) return f.stream.Read(b)
} }
func (f *fileStreamOut) EnvFile(fs filestore.FileStore) (envexec.File, error) { func (f *fileStreamOut) EnvFile(fs filestore.FileStore) (envexec.File, error) {

View File

@ -7,6 +7,7 @@ import (
"io" "io"
"github.com/criyle/go-judge/cmd/go-judge/model" "github.com/criyle/go-judge/cmd/go-judge/model"
"github.com/criyle/go-judge/envexec"
"github.com/criyle/go-judge/worker" "github.com/criyle/go-judge/worker"
"go.uber.org/zap" "go.uber.org/zap"
"golang.org/x/sync/errgroup" "golang.org/x/sync/errgroup"
@ -242,11 +243,12 @@ func streamInput(ctx context.Context, s Stream, si []*fileStreamIn, execCancel f
if !ok { if !ok {
return fmt.Errorf("input does not exist: %d/%d", in.Resize.Index, in.Resize.Fd) return fmt.Errorf("input does not exist: %d/%d", in.Resize.Index, in.Resize.Fd)
} }
tty := f.GetTTY() if err = f.SetSize(&envexec.TerminalSize{
if tty == nil { Cols: uint16(in.Resize.Cols),
return fmt.Errorf("resize input does not have tty: %d/%d", in.Resize.Index, in.Resize.Fd) Rows: uint16(in.Resize.Rows),
} X: uint16(in.Resize.X),
if err = setWinsize(tty, in.Resize); err != nil { Y: uint16(in.Resize.Y),
}); err != nil {
return fmt.Errorf("resize %d/%d: %w", in.Resize.Index, in.Resize.Fd, err) return fmt.Errorf("resize %d/%d: %w", in.Resize.Index, in.Resize.Fd, err)
} }

View File

@ -1,11 +0,0 @@
//go:build !linux && !darwin
package stream
import (
"os"
)
func setWinsize(f *os.File, i *ResizeRequest) error {
return nil
}

View File

@ -1,19 +0,0 @@
//go:build linux || darwin
package stream
import (
"os"
"github.com/creack/pty"
)
func setWinsize(f *os.File, i *ResizeRequest) error {
winSize := &pty.Winsize{
Rows: uint16(i.Rows),
Cols: uint16(i.Cols),
X: uint16(i.X),
Y: uint16(i.Y),
}
return pty.Setsize(f, winSize)
}

View File

@ -5,6 +5,8 @@ import (
"io" "io"
"os" "os"
"sync" "sync"
"github.com/creack/pty"
) )
var ( var (
@ -42,88 +44,142 @@ func NewFileReader(r io.Reader) File {
// responsibility to close the WritePipe // responsibility to close the WritePipe
type FileStreamIn interface { type FileStreamIn interface {
File File
Done() <-chan struct{} io.WriteCloser
WritePipe() *os.File SetSize(*TerminalSize) error
Close() error }
// TerminalSize controls the size of the terminal if TTY is enabled
type TerminalSize struct {
Rows uint16 // ws_row: Number of rows (in cells).
Cols uint16 // ws_col: Number of columns (in cells).
X uint16 // ws_xpixel: Width in pixels.
Y uint16 // ws_ypixel: Height in pixels.
} }
type fileStreamIn struct { type fileStreamIn struct {
done chan struct{} started chan struct{}
closed chan struct{}
w *sharedFile w *sharedFile
} }
func NewFileStreamIn() FileStreamIn { func NewFileStreamIn() FileStreamIn {
return &fileStreamIn{ return &fileStreamIn{
done: make(chan struct{}), started: make(chan struct{}),
closed: make(chan struct{}),
} }
} }
func (*fileStreamIn) isFile() {} func (*fileStreamIn) isFile() {}
func (f *fileStreamIn) start(w *sharedFile) { func (f *fileStreamIn) start(w *sharedFile) {
select {
case <-f.closed:
w.Close()
default:
f.w = w f.w = w
close(f.done) close(f.started)
}
} }
func (f *fileStreamIn) Done() <-chan struct{} { func (f *fileStreamIn) Write(b []byte) (int, error) {
return f.done select {
case <-f.started:
return f.w.f.Write(b)
case <-f.closed:
return 0, io.EOF
}
} }
func (f *fileStreamIn) WritePipe() *os.File { func (f *fileStreamIn) SetSize(s *TerminalSize) error {
<-f.done select {
return f.w.f case <-f.started:
return pty.Setsize(f.w.f, &pty.Winsize{
Rows: s.Rows,
Cols: s.Cols,
X: s.X,
Y: s.Y,
})
case <-f.closed:
return io.ErrClosedPipe
}
} }
func (f *fileStreamIn) Close() error { func (f *fileStreamIn) Close() error {
if f.w != nil { select {
case <-f.started:
close(f.closed)
return f.w.Close() return f.w.Close()
}
case <-f.closed:
return io.ErrClosedPipe
default:
close(f.closed)
return nil return nil
} }
}
// FileStreamOut represent a out streaming pipe and the streamer is able to read // FileStreamOut represent a out streaming pipe and the streamer is able to read
// to the read end of the pipe after pipe created. It is the callers // to the read end of the pipe after pipe created. It is the callers
// responsibility to close the ReadPipe // responsibility to close the ReadPipe
type FileStreamOut interface { type FileStreamOut interface {
File File
Done() <-chan struct{} io.ReadCloser
ReadPipe() *os.File
Close() error
} }
type fileStreamOut struct { type fileStreamOut struct {
done chan struct{} started chan struct{}
closed chan struct{}
r *sharedFile r *sharedFile
} }
func NewFileStreamOut() FileStreamOut { func NewFileStreamOut() FileStreamOut {
return &fileStreamOut{ return &fileStreamOut{
done: make(chan struct{}), started: make(chan struct{}),
closed: make(chan struct{}),
} }
} }
func (*fileStreamOut) isFile() {} func (*fileStreamOut) isFile() {}
func (f *fileStreamOut) start(r *sharedFile) { func (f *fileStreamOut) start(r *sharedFile) {
select {
case <-f.closed:
r.Close()
default:
f.r = r f.r = r
close(f.done) close(f.started)
}
} }
func (f *fileStreamOut) Done() <-chan struct{} { func (f *fileStreamOut) Read(p []byte) (n int, err error) {
return f.done select {
} case <-f.started:
return f.r.f.Read(p)
func (f *fileStreamOut) ReadPipe() *os.File { case <-f.closed:
<-f.done return 0, io.EOF
return f.r.f }
} }
func (f *fileStreamOut) Close() error { func (f *fileStreamOut) Close() error {
if f.r != nil { select {
case <-f.started:
close(f.closed)
return f.r.Close() return f.r.Close()
}
case <-f.closed:
return io.ErrClosedPipe
default:
close(f.closed)
return nil return nil
} }
}
// FileInput represent file input which will be opened in read-only mode // FileInput represent file input which will be opened in read-only mode
type FileInput struct { type FileInput struct {
@ -204,7 +260,7 @@ type sharedFile struct {
count int count int
} }
func newShreadFile(f *os.File) *sharedFile { func newSharedFile(f *os.File) *sharedFile {
return &sharedFile{f: f, count: 0} return &sharedFile{f: f, count: 0}
} }

View File

@ -31,7 +31,7 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
err = fmt.Errorf("failed to open tty %v", err) err = fmt.Errorf("failed to open tty %v", err)
return nil, nil, err return nil, nil, err
} }
sf := newShreadFile(fPty) sf := newSharedFile(fPty)
files := make([]*os.File, count) files := make([]*os.File, count)
pipeToCollect := make([]pipeCollector, 0) pipeToCollect := make([]pipeCollector, 0)
@ -145,6 +145,34 @@ func prepareCmdFdNoTty(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.Fi
}() }()
// record the same file to avoid multiple file open // record the same file to avoid multiple file open
cf := make(map[string]*os.File) cf := make(map[string]*os.File)
prepareFileCollector := func(t *FileCollector) (*os.File, error) {
if f, ok := cf[t.Name]; ok {
return f, nil
}
if t.Pipe {
b, err := newPipeBuffer(t.Limit, newFileStore)
if err != nil {
return nil, fmt.Errorf("pipe: create: %w", err)
}
cf[t.Name] = b.W
pipeToCollect = append(pipeToCollect, pipeCollector{b.Done, b.Buffer, t.Limit, t.Name, true})
return b.W, nil
} else {
f, err := c.Environment.Open(t.Name, os.O_CREATE|os.O_WRONLY, 0777)
if err != nil {
return nil, fmt.Errorf("container: create file %v: %w", t.Name, err)
}
buffer, err := c.Environment.Open(t.Name, os.O_RDWR, 0777)
if err != nil {
f.Close()
return nil, fmt.Errorf("container: open created file %v: %w", t.Name, err)
}
cf[t.Name] = f
pipeToCollect = append(pipeToCollect, pipeCollector{closedChan, buffer, t.Limit, t.Name, false})
return f, nil
}
}
for j, t := range c.Files { for j, t := range c.Files {
switch t := t.(type) { switch t := t.(type) {
@ -167,35 +195,11 @@ func prepareCmdFdNoTty(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.Fi
files[j] = f files[j] = f
case *FileCollector: case *FileCollector:
if f, ok := cf[t.Name]; ok { f, err := prepareFileCollector(t)
if err != nil {
return nil, nil, err
}
files[j] = f files[j] = f
break
}
if t.Pipe {
b, err := newPipeBuffer(t.Limit, newFileStore)
if err != nil {
return nil, nil, fmt.Errorf("pipe: create: %w", err)
}
cf[t.Name] = b.W
files[j] = b.W
pipeToCollect = append(pipeToCollect, pipeCollector{b.Done, b.Buffer, t.Limit, t.Name, true})
} else {
f, err := c.Environment.Open(t.Name, os.O_CREATE|os.O_WRONLY, 0777)
if err != nil {
return nil, nil, fmt.Errorf("container: create file %v: %w", t.Name, err)
}
cf[t.Name] = f
buffer, err := c.Environment.Open(t.Name, os.O_RDWR, 0777)
if err != nil {
return nil, nil, fmt.Errorf("container: open created file %v: %w", t.Name, err)
}
files[j] = f
pipeToCollect = append(pipeToCollect, pipeCollector{closedChan, buffer, t.Limit, t.Name, false})
}
case *FileWriter: case *FileWriter:
_, w, err := newPipe(t.Writer, t.Limit) _, w, err := newPipe(t.Writer, t.Limit)