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

This commit is contained in:
criyle 2025-05-29 04:40:53 -04:00
parent 033790c1c0
commit c5575ffe2b
5 changed files with 189 additions and 87 deletions

View File

@ -2,8 +2,6 @@ package stream
import ( import (
"fmt" "fmt"
"io"
"math"
"os" "os"
"github.com/criyle/go-judge/envexec" "github.com/criyle/go-judge/envexec"
@ -14,50 +12,37 @@ import (
var ( var (
_ worker.CmdFile = &fileStreamIn{} _ worker.CmdFile = &fileStreamIn{}
_ worker.CmdFile = &fileStreamOut{} _ worker.CmdFile = &fileStreamOut{}
_ envexec.ReaderTTY = &fileStreamInReader{}
) )
type fileStreamIn struct { type fileStreamIn struct {
stream envexec.FileStreamIn
index int index int
fd int fd int
r io.ReadCloser
w *io.PipeWriter
tty *os.File
done chan struct{}
hasTTY bool 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 { func newFileStreamIn(index, fd int, hasTTY bool) *fileStreamIn {
r, w := io.Pipe() return &fileStreamIn{
fi := &fileStreamIn{index: index, fd: fd, w: w, done: make(chan struct{}), hasTTY: hasTTY} stream: envexec.NewFileStreamIn(),
fi.r = &fileStreamInReader{r, fi} index: index,
return fi fd: fd,
hasTTY: hasTTY,
}
} }
func (f *fileStreamIn) GetTTY() *os.File { func (f *fileStreamIn) GetTTY() *os.File {
if !f.hasTTY { if !f.hasTTY {
return nil return nil
} }
<-f.done return f.stream.WritePipe()
return f.tty
} }
func (f *fileStreamIn) Write(b []byte) (int, error) { 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) { 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 { func (f *fileStreamIn) String() string {
@ -65,28 +50,29 @@ func (f *fileStreamIn) String() string {
} }
func (f *fileStreamIn) Close() error { func (f *fileStreamIn) Close() error {
f.r.Close() return f.stream.Close()
return f.w.Close()
} }
type fileStreamOut struct { type fileStreamOut struct {
stream envexec.FileStreamOut
index int index int
fd int fd int
r *io.PipeReader
w *io.PipeWriter
} }
func newFileStreamOut(index, fd int) *fileStreamOut { func newFileStreamOut(index, fd int) *fileStreamOut {
r, w := io.Pipe() return &fileStreamOut{
return &fileStreamOut{index: index, fd: fd, r: r, w: w} stream: envexec.NewFileStreamOut(),
index: index,
fd: fd,
}
} }
func (f *fileStreamOut) Read(b []byte) (int, error) { 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) { 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 { func (f *fileStreamOut) String() string {
@ -94,6 +80,5 @@ func (f *fileStreamOut) String() string {
} }
func (f *fileStreamOut) Close() error { func (f *fileStreamOut) Close() error {
f.w.Close() return f.stream.Close()
return f.r.Close()
} }

View File

@ -127,7 +127,6 @@ func Start(baseCtx context.Context, s Stream, w worker.Worker, srcPrefix []strin
cancel() cancel()
closeFunc() closeFunc()
streamOut = nil
wg.Wait() wg.Wait()
return err return err
} }

View File

@ -4,6 +4,17 @@ import (
"fmt" "fmt"
"io" "io"
"os" "os"
"sync"
)
var (
_ File = &FileReader{}
_ File = &fileStreamIn{}
_ File = &fileStreamOut{}
_ File = &FileInput{}
_ File = &FileCollector{}
_ File = &FileWriter{}
_ File = &FileOpened{}
) )
// File defines interface of envexec files // File defines interface of envexec files
@ -20,16 +31,98 @@ type FileReader struct {
func (*FileReader) isFile() {} func (*FileReader) isFile() {}
// NewFileReader creates File input which can be fully read before exec // NewFileReader creates File input which can be fully read before exec.
// or piped into exec // If pipe is required, use the FileStream to get the write end of pipe instead
func NewFileReader(r io.Reader, s bool) File { func NewFileReader(r io.Reader) File {
return &FileReader{Reader: r, Stream: s} return &FileReader{Reader: r}
} }
// ReaderTTY will be asserts when File Reader is provided and TTY is enabled // FileStreamIn represent a input streaming pipe and the streamer is able to write
// and then TTY will be called with pty file // to the write end of the pipe after pipe created. It is the callers
type ReaderTTY interface { // responsibility to close the WritePipe
TTY(*os.File) 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 // 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) 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
}

View File

@ -4,7 +4,6 @@ import (
"fmt" "fmt"
"io" "io"
"os" "os"
"sync"
"github.com/creack/pty" "github.com/creack/pty"
) )
@ -16,9 +15,15 @@ func init() {
close(closedChan) 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 // prepare Files for tty input / output
func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File, p []pipeCollector, err error) { func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File, p []pipeCollector, err error) {
var wg sync.WaitGroup
var hasInput, hasOutput bool var hasInput, hasOutput bool
fPty, fTty, err := pty.Open() 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) err = fmt.Errorf("failed to open tty %v", err)
return nil, nil, err return nil, nil, err
} }
sf := newShreadFile(fPty)
files := make([]*os.File, count) files := make([]*os.File, count)
pipeToCollect := make([]pipeCollector, 0) pipeToCollect := make([]pipeCollector, 0)
@ -34,7 +40,6 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
if err != nil { if err != nil {
closeFiles(files...) closeFiles(files...)
closeFiles(fTty, fPty) closeFiles(fTty, fPty)
wg.Wait()
} }
}() }()
@ -44,6 +49,24 @@ func prepareCmdFdTTY(c *Cmd, count int, newStoreFile NewStoreFile) (f []*os.File
case *FileOpened: case *FileOpened:
files[j] = t.File 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: case *FileReader:
if hasInput { if hasInput {
return nil, nil, fmt.Errorf("tty: multiple input not allowed") 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 files[j] = fTty
// copy input // copy input
wg.Add(1) sf.Acquire()
go func() { go func() {
defer wg.Done() defer sf.Close()
io.Copy(fPty, t.Reader) io.Copy(fPty, t.Reader)
}() }()
// provide TTY
if tty, ok := t.Reader.(ReaderTTY); ok {
tty.TTY(fPty)
}
case *FileInput: case *FileInput:
var f *os.File var f *os.File
f, err = os.Open(t.Path) 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}) pipeToCollect = append(pipeToCollect, pipeCollector{done, buf, t.Limit, t.Name, true})
wg.Add(1) sf.Acquire()
go func() { go func() {
defer close(done) defer close(done)
defer wg.Done() defer sf.Close()
io.CopyN(buf, fPty, int64(t.Limit)+1) 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 hasOutput = true
wg.Add(1) sf.Acquire()
go func() { go func() {
defer wg.Done() defer sf.Close()
io.Copy(t.Writer, fPty) 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) 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 return files, pipeToCollect, nil
} }
func prepareCmdFd(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p []pipeCollector, err error) { func prepareCmdFdNoTty(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p []pipeCollector, err error) {
if c.TTY {
return prepareCmdFdTTY(c, count, newFileStore)
}
files := make([]*os.File, count) files := make([]*os.File, count)
pipeToCollect := make([]pipeCollector, 0) pipeToCollect := make([]pipeCollector, 0)
defer func() { defer func() {
@ -144,21 +153,11 @@ func prepareCmdFd(c *Cmd, count int, newFileStore NewStoreFile) (f []*os.File, p
files[j] = t.File files[j] = t.File
case *FileReader: 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) f, err := readerToFile(t.Reader)
if err != nil { if err != nil {
return nil, nil, fmt.Errorf("pipe: open reader: %w", err) return nil, nil, fmt.Errorf("pipe: open reader: %w", err)
} }
files[j] = f files[j] = f
}
case *FileInput: case *FileInput:
f, err := os.Open(t.Path) f, err := os.Open(t.Path)

View File

@ -44,7 +44,7 @@ type MemoryFile struct {
// EnvFile prepares file for envexec file // EnvFile prepares file for envexec file
func (f *MemoryFile) EnvFile(fs filestore.FileStore) (envexec.File, error) { 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 { func (f *MemoryFile) String() string {