mirror of
https://github.com/criyle/go-judge.git
synced 2025-09-26 22:39:12 +08:00
refactor(envexec): support file stream in & out directly
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
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:
parent
033790c1c0
commit
c5575ffe2b
@ -2,8 +2,6 @@ package stream
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"os"
|
||||
|
||||
"github.com/criyle/go-judge/envexec"
|
||||
@ -14,50 +12,37 @@ import (
|
||||
var (
|
||||
_ worker.CmdFile = &fileStreamIn{}
|
||||
_ worker.CmdFile = &fileStreamOut{}
|
||||
_ envexec.ReaderTTY = &fileStreamInReader{}
|
||||
)
|
||||
|
||||
type fileStreamIn struct {
|
||||
stream envexec.FileStreamIn
|
||||
index int
|
||||
fd int
|
||||
r io.ReadCloser
|
||||
w *io.PipeWriter
|
||||
tty *os.File
|
||||
done chan struct{}
|
||||
hasTTY bool
|
||||
}
|
||||
|
||||
type fileStreamInReader struct {
|
||||
*io.PipeReader
|
||||
fi *fileStreamIn
|
||||
}
|
||||
|
||||
func (f *fileStreamInReader) TTY(tty *os.File) {
|
||||
f.fi.tty = tty
|
||||
close(f.fi.done)
|
||||
}
|
||||
|
||||
func newFileStreamIn(index, fd int, hasTTY bool) *fileStreamIn {
|
||||
r, w := io.Pipe()
|
||||
fi := &fileStreamIn{index: index, fd: fd, w: w, done: make(chan struct{}), hasTTY: hasTTY}
|
||||
fi.r = &fileStreamInReader{r, fi}
|
||||
return fi
|
||||
return &fileStreamIn{
|
||||
stream: envexec.NewFileStreamIn(),
|
||||
index: index,
|
||||
fd: fd,
|
||||
hasTTY: hasTTY,
|
||||
}
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) GetTTY() *os.File {
|
||||
if !f.hasTTY {
|
||||
return nil
|
||||
}
|
||||
<-f.done
|
||||
return f.tty
|
||||
return f.stream.WritePipe()
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) Write(b []byte) (int, error) {
|
||||
return f.w.Write(b)
|
||||
return f.stream.WritePipe().Write(b)
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) EnvFile(fs filestore.FileStore) (envexec.File, error) {
|
||||
return envexec.NewFileReader(f.r, true), nil
|
||||
return f.stream, nil
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) String() string {
|
||||
@ -65,28 +50,29 @@ func (f *fileStreamIn) String() string {
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) Close() error {
|
||||
f.r.Close()
|
||||
return f.w.Close()
|
||||
return f.stream.Close()
|
||||
}
|
||||
|
||||
type fileStreamOut struct {
|
||||
stream envexec.FileStreamOut
|
||||
index int
|
||||
fd int
|
||||
r *io.PipeReader
|
||||
w *io.PipeWriter
|
||||
}
|
||||
|
||||
func newFileStreamOut(index, fd int) *fileStreamOut {
|
||||
r, w := io.Pipe()
|
||||
return &fileStreamOut{index: index, fd: fd, r: r, w: w}
|
||||
return &fileStreamOut{
|
||||
stream: envexec.NewFileStreamOut(),
|
||||
index: index,
|
||||
fd: fd,
|
||||
}
|
||||
}
|
||||
|
||||
func (f *fileStreamOut) Read(b []byte) (int, error) {
|
||||
return f.r.Read(b)
|
||||
return f.stream.ReadPipe().Read(b)
|
||||
}
|
||||
|
||||
func (f *fileStreamOut) EnvFile(fs filestore.FileStore) (envexec.File, error) {
|
||||
return envexec.NewFileWriter(f.w, envexec.Size(math.MaxInt32)), nil
|
||||
return f.stream, nil
|
||||
}
|
||||
|
||||
func (f *fileStreamOut) String() string {
|
||||
@ -94,6 +80,5 @@ func (f *fileStreamOut) String() string {
|
||||
}
|
||||
|
||||
func (f *fileStreamOut) Close() error {
|
||||
f.w.Close()
|
||||
return f.r.Close()
|
||||
return f.stream.Close()
|
||||
}
|
||||
|
@ -127,7 +127,6 @@ func Start(baseCtx context.Context, s Stream, w worker.Worker, srcPrefix []strin
|
||||
|
||||
cancel()
|
||||
closeFunc()
|
||||
streamOut = nil
|
||||
wg.Wait()
|
||||
return err
|
||||
}
|
||||
|
135
envexec/file.go
135
envexec/file.go
@ -4,6 +4,17 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var (
|
||||
_ File = &FileReader{}
|
||||
_ File = &fileStreamIn{}
|
||||
_ File = &fileStreamOut{}
|
||||
_ File = &FileInput{}
|
||||
_ File = &FileCollector{}
|
||||
_ File = &FileWriter{}
|
||||
_ File = &FileOpened{}
|
||||
)
|
||||
|
||||
// File defines interface of envexec files
|
||||
@ -20,16 +31,98 @@ type FileReader struct {
|
||||
|
||||
func (*FileReader) isFile() {}
|
||||
|
||||
// NewFileReader creates File input which can be fully read before exec
|
||||
// or piped into exec
|
||||
func NewFileReader(r io.Reader, s bool) File {
|
||||
return &FileReader{Reader: r, Stream: s}
|
||||
// NewFileReader creates File input which can be fully read before exec.
|
||||
// If pipe is required, use the FileStream to get the write end of pipe instead
|
||||
func NewFileReader(r io.Reader) File {
|
||||
return &FileReader{Reader: r}
|
||||
}
|
||||
|
||||
// ReaderTTY will be asserts when File Reader is provided and TTY is enabled
|
||||
// and then TTY will be called with pty file
|
||||
type ReaderTTY interface {
|
||||
TTY(*os.File)
|
||||
// FileStreamIn represent a input streaming pipe and the streamer is able to write
|
||||
// to the write end of the pipe after pipe created. It is the callers
|
||||
// responsibility to close the WritePipe
|
||||
type FileStreamIn interface {
|
||||
File
|
||||
Done() <-chan struct{}
|
||||
WritePipe() *os.File
|
||||
Close() error
|
||||
}
|
||||
|
||||
type fileStreamIn struct {
|
||||
done chan struct{}
|
||||
w *sharedFile
|
||||
}
|
||||
|
||||
func NewFileStreamIn() FileStreamIn {
|
||||
return &fileStreamIn{
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (*fileStreamIn) isFile() {}
|
||||
|
||||
func (f *fileStreamIn) start(w *sharedFile) {
|
||||
f.w = w
|
||||
close(f.done)
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) Done() <-chan struct{} {
|
||||
return f.done
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) WritePipe() *os.File {
|
||||
<-f.done
|
||||
return f.w.f
|
||||
}
|
||||
|
||||
func (f *fileStreamIn) Close() error {
|
||||
if f.w != nil {
|
||||
return f.w.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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
|
||||
// responsibility to close the ReadPipe
|
||||
type FileStreamOut interface {
|
||||
File
|
||||
Done() <-chan struct{}
|
||||
ReadPipe() *os.File
|
||||
Close() error
|
||||
}
|
||||
|
||||
type fileStreamOut struct {
|
||||
done chan struct{}
|
||||
r *sharedFile
|
||||
}
|
||||
|
||||
func NewFileStreamOut() FileStreamOut {
|
||||
return &fileStreamOut{
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (*fileStreamOut) isFile() {}
|
||||
|
||||
func (f *fileStreamOut) start(r *sharedFile) {
|
||||
f.r = r
|
||||
close(f.done)
|
||||
}
|
||||
|
||||
func (f *fileStreamOut) Done() <-chan struct{} {
|
||||
return f.done
|
||||
}
|
||||
|
||||
func (f *fileStreamOut) ReadPipe() *os.File {
|
||||
<-f.done
|
||||
return f.r.f
|
||||
}
|
||||
|
||||
func (f *fileStreamOut) Close() error {
|
||||
if f.r != nil {
|
||||
return f.r.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// FileInput represent file input which will be opened in read-only mode
|
||||
@ -104,3 +197,29 @@ func FileToReader(f File) (io.ReadCloser, error) {
|
||||
return nil, fmt.Errorf("file cannot open as reader: %T", f)
|
||||
}
|
||||
}
|
||||
|
||||
type sharedFile struct {
|
||||
mu sync.Mutex
|
||||
f *os.File
|
||||
count int
|
||||
}
|
||||
|
||||
func newShreadFile(f *os.File) *sharedFile {
|
||||
return &sharedFile{f: f, count: 0}
|
||||
}
|
||||
|
||||
func (s *sharedFile) Acquire() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.count++
|
||||
}
|
||||
|
||||
func (s *sharedFile) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.count--
|
||||
if s.count == 0 {
|
||||
return s.f.Close()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
@ -4,7 +4,6 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"github.com/creack/pty"
|
||||
)
|
||||
@ -16,9 +15,15 @@ func init() {
|
||||
close(closedChan)
|
||||
}
|
||||
|
||||
func prepareCmdFd(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p []pipeCollector, err error) {
|
||||
if c.TTY {
|
||||
return prepareCmdFdTTY(c, count, newFileStore)
|
||||
}
|
||||
return prepareCmdFdNoTty(c, count, newFileStore)
|
||||
}
|
||||
|
||||
// prepare Files for tty input / output
|
||||
func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File, p []pipeCollector, err error) {
|
||||
var wg sync.WaitGroup
|
||||
var hasInput, hasOutput bool
|
||||
|
||||
fPty, fTty, err := pty.Open()
|
||||
@ -26,6 +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)
|
||||
|
||||
files := make([]*os.File, count)
|
||||
pipeToCollect := make([]pipeCollector, 0)
|
||||
@ -34,7 +40,6 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
|
||||
if err != nil {
|
||||
closeFiles(files...)
|
||||
closeFiles(fTty, fPty)
|
||||
wg.Wait()
|
||||
}
|
||||
}()
|
||||
|
||||
@ -44,6 +49,24 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
|
||||
case *FileOpened:
|
||||
files[j] = t.File
|
||||
|
||||
case *fileStreamIn:
|
||||
if hasInput {
|
||||
return nil, nil, fmt.Errorf("tty: multiple input not allowed")
|
||||
}
|
||||
hasInput = true
|
||||
|
||||
files[j] = fTty
|
||||
// provide TTY
|
||||
sf.Acquire()
|
||||
t.start(sf)
|
||||
|
||||
case *fileStreamOut:
|
||||
files[j] = fTty
|
||||
hasOutput = true
|
||||
// Provide TTY
|
||||
sf.Acquire()
|
||||
t.start(sf)
|
||||
|
||||
case *FileReader:
|
||||
if hasInput {
|
||||
return nil, nil, fmt.Errorf("tty: multiple input not allowed")
|
||||
@ -53,17 +76,12 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
|
||||
files[j] = fTty
|
||||
|
||||
// copy input
|
||||
wg.Add(1)
|
||||
sf.Acquire()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer sf.Close()
|
||||
io.Copy(fPty, t.Reader)
|
||||
}()
|
||||
|
||||
// provide TTY
|
||||
if tty, ok := t.Reader.(ReaderTTY); ok {
|
||||
tty.TTY(fPty)
|
||||
}
|
||||
|
||||
case *FileInput:
|
||||
var f *os.File
|
||||
f, err = os.Open(t.Path)
|
||||
@ -86,10 +104,10 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
|
||||
}
|
||||
pipeToCollect = append(pipeToCollect, pipeCollector{done, buf, t.Limit, t.Name, true})
|
||||
|
||||
wg.Add(1)
|
||||
sf.Acquire()
|
||||
go func() {
|
||||
defer close(done)
|
||||
defer wg.Done()
|
||||
defer sf.Close()
|
||||
io.CopyN(buf, fPty, int64(t.Limit)+1)
|
||||
}()
|
||||
|
||||
@ -100,9 +118,9 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
|
||||
}
|
||||
hasOutput = true
|
||||
|
||||
wg.Add(1)
|
||||
sf.Acquire()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer sf.Close()
|
||||
io.Copy(t.Writer, fPty)
|
||||
}()
|
||||
|
||||
@ -110,19 +128,10 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
|
||||
return nil, nil, fmt.Errorf("tty: unknown file type: %T", t)
|
||||
}
|
||||
}
|
||||
|
||||
// ensure pty close after use
|
||||
go func() {
|
||||
wg.Wait()
|
||||
fPty.Close()
|
||||
}()
|
||||
return files, pipeToCollect, nil
|
||||
}
|
||||
|
||||
func prepareCmdFd(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p []pipeCollector, err error) {
|
||||
if c.TTY {
|
||||
return prepareCmdFdTTY(c, count, newFileStore)
|
||||
}
|
||||
func prepareCmdFdNoTty(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p []pipeCollector, err error) {
|
||||
files := make([]*os.File, count)
|
||||
pipeToCollect := make([]pipeCollector, 0)
|
||||
defer func() {
|
||||
@ -144,21 +153,11 @@ func prepareCmdFd(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p
|
||||
files[j] = t.File
|
||||
|
||||
case *FileReader:
|
||||
if t.Stream {
|
||||
r, w, err := os.Pipe()
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("pipe: create: %w", err)
|
||||
}
|
||||
go w.ReadFrom(t.Reader)
|
||||
|
||||
files[j] = r
|
||||
} else {
|
||||
f, err := readerToFile(t.Reader)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("pipe: open reader: %w", err)
|
||||
}
|
||||
files[j] = f
|
||||
}
|
||||
|
||||
case *FileInput:
|
||||
f, err := os.Open(t.Path)
|
||||
|
@ -44,7 +44,7 @@ type MemoryFile struct {
|
||||
|
||||
// EnvFile prepares file for envexec file
|
||||
func (f *MemoryFile) EnvFile(fs filestore.FileStore) (envexec.File, error) {
|
||||
return envexec.NewFileReader(bytes.NewReader(f.Content), false), nil
|
||||
return envexec.NewFileReader(bytes.NewReader(f.Content)), nil
|
||||
}
|
||||
|
||||
func (f *MemoryFile) String() string {
|
||||
|
Loading…
Reference in New Issue
Block a user