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
import (
"errors"
"fmt"
"os"
"github.com/criyle/go-judge/envexec"
"github.com/criyle/go-judge/filestore"
@ -12,6 +12,8 @@ import (
var (
_ worker.CmdFile = &fileStreamIn{}
_ worker.CmdFile = &fileStreamOut{}
errNoTTY = errors.New("stream is not a tty")
)
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 {
return nil
return errNoTTY
}
return f.stream.WritePipe()
return f.stream.SetSize(s)
}
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) {
@ -68,7 +70,7 @@ func newFileStreamOut(index, fd int) *fileStreamOut {
}
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) {

View File

@ -7,6 +7,7 @@ import (
"io"
"github.com/criyle/go-judge/cmd/go-judge/model"
"github.com/criyle/go-judge/envexec"
"github.com/criyle/go-judge/worker"
"go.uber.org/zap"
"golang.org/x/sync/errgroup"
@ -242,11 +243,12 @@ func streamInput(ctx context.Context, s Stream, si []*fileStreamIn, execCancel f
if !ok {
return fmt.Errorf("input does not exist: %d/%d", in.Resize.Index, in.Resize.Fd)
}
tty := f.GetTTY()
if tty == nil {
return fmt.Errorf("resize input does not have tty: %d/%d", in.Resize.Index, in.Resize.Fd)
}
if err = setWinsize(tty, in.Resize); err != nil {
if err = f.SetSize(&envexec.TerminalSize{
Cols: uint16(in.Resize.Cols),
Rows: uint16(in.Resize.Rows),
X: uint16(in.Resize.X),
Y: uint16(in.Resize.Y),
}); err != nil {
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"
"os"
"sync"
"github.com/creack/pty"
)
var (
@ -42,43 +44,82 @@ func NewFileReader(r io.Reader) File {
// responsibility to close the WritePipe
type FileStreamIn interface {
File
Done() <-chan struct{}
WritePipe() *os.File
Close() error
io.WriteCloser
SetSize(*TerminalSize) 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 {
done chan struct{}
w *sharedFile
started chan struct{}
closed chan struct{}
w *sharedFile
}
func NewFileStreamIn() FileStreamIn {
return &fileStreamIn{
done: make(chan struct{}),
started: make(chan struct{}),
closed: make(chan struct{}),
}
}
func (*fileStreamIn) isFile() {}
func (f *fileStreamIn) start(w *sharedFile) {
f.w = w
close(f.done)
select {
case <-f.closed:
w.Close()
default:
f.w = w
close(f.started)
}
}
func (f *fileStreamIn) Done() <-chan struct{} {
return f.done
func (f *fileStreamIn) Write(b []byte) (int, error) {
select {
case <-f.started:
return f.w.f.Write(b)
case <-f.closed:
return 0, io.EOF
}
}
func (f *fileStreamIn) WritePipe() *os.File {
<-f.done
return f.w.f
func (f *fileStreamIn) SetSize(s *TerminalSize) error {
select {
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 {
if f.w != nil {
select {
case <-f.started:
close(f.closed)
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
@ -86,43 +127,58 @@ func (f *fileStreamIn) Close() error {
// responsibility to close the ReadPipe
type FileStreamOut interface {
File
Done() <-chan struct{}
ReadPipe() *os.File
Close() error
io.ReadCloser
}
type fileStreamOut struct {
done chan struct{}
r *sharedFile
started chan struct{}
closed chan struct{}
r *sharedFile
}
func NewFileStreamOut() FileStreamOut {
return &fileStreamOut{
done: make(chan struct{}),
started: make(chan struct{}),
closed: make(chan struct{}),
}
}
func (*fileStreamOut) isFile() {}
func (f *fileStreamOut) start(r *sharedFile) {
f.r = r
close(f.done)
select {
case <-f.closed:
r.Close()
default:
f.r = r
close(f.started)
}
}
func (f *fileStreamOut) Done() <-chan struct{} {
return f.done
}
func (f *fileStreamOut) Read(p []byte) (n int, err error) {
select {
case <-f.started:
return f.r.f.Read(p)
func (f *fileStreamOut) ReadPipe() *os.File {
<-f.done
return f.r.f
case <-f.closed:
return 0, io.EOF
}
}
func (f *fileStreamOut) Close() error {
if f.r != nil {
select {
case <-f.started:
close(f.closed)
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
@ -204,7 +260,7 @@ type sharedFile struct {
count int
}
func newShreadFile(f *os.File) *sharedFile {
func newSharedFile(f *os.File) *sharedFile {
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)
return nil, nil, err
}
sf := newShreadFile(fPty)
sf := newSharedFile(fPty)
files := make([]*os.File, count)
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
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 {
switch t := t.(type) {
@ -167,35 +195,11 @@ func prepareCmdFdNoTty(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.Fi
files[j] = f
case *FileCollector:
if f, ok := cf[t.Name]; ok {
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})
f, err := prepareFileCollector(t)
if err != nil {
return nil, nil, err
}
files[j] = f
case *FileWriter:
_, w, err := newPipe(t.Writer, t.Limit)