diff --git a/cmd/go-judge/stream/file.go b/cmd/go-judge/stream/file.go index ddd89ae..22b9a5e 100644 --- a/cmd/go-judge/stream/file.go +++ b/cmd/go-judge/stream/file.go @@ -2,8 +2,6 @@ package stream import ( "fmt" - "io" - "math" "os" "github.com/criyle/go-judge/envexec" @@ -12,52 +10,39 @@ import ( ) var ( - _ worker.CmdFile = &fileStreamIn{} - _ worker.CmdFile = &fileStreamOut{} - _ envexec.ReaderTTY = &fileStreamInReader{} + _ worker.CmdFile = &fileStreamIn{} + _ worker.CmdFile = &fileStreamOut{} ) 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 { - index int - fd int - r *io.PipeReader - w *io.PipeWriter + stream envexec.FileStreamOut + index int + fd int } 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() } diff --git a/cmd/go-judge/stream/stream.go b/cmd/go-judge/stream/stream.go index 4b3bcd5..77b7c93 100644 --- a/cmd/go-judge/stream/stream.go +++ b/cmd/go-judge/stream/stream.go @@ -127,7 +127,6 @@ func Start(baseCtx context.Context, s Stream, w worker.Worker, srcPrefix []strin cancel() closeFunc() - streamOut = nil wg.Wait() return err } diff --git a/envexec/file.go b/envexec/file.go index 5b8a56e..38329ee 100644 --- a/envexec/file.go +++ b/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 +} diff --git a/envexec/file_prepare.go b/envexec/file_prepare.go index 787ef20..e64d8e7 100644 --- a/envexec/file_prepare.go +++ b/envexec/file_prepare.go @@ -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 + 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) diff --git a/worker/file.go b/worker/file.go index c17819d..a82b08d 100644 --- a/worker/file.go +++ b/worker/file.go @@ -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 {