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 (
"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()
}

View File

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

View File

@ -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
}

View File

@ -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)

View File

@ -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 {