refactor(envexec): support file stream in & out directly

Author: criyle

cmd/go-judge/stream/file.go

@@ -2,8 +2,6 @@ package stream
 
 import (
 	"fmt"
-	"io"
-	"math"
 	"os"
 
 	"github.com/criyle/go-judge/envexec"
@@ -12,88 +10,75 @@ 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 {
 	return fmt.Sprintf("fileStreamIn:(index:%d,fd:%d)", f.index, f.fd)
 }
 
 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 {
 	return fmt.Sprintf("fileStreamOut:(index:%d,fd:%d)", f.index, f.fd)
 }
 
 func (f *fileStreamOut) Close() error {
-	f.w.Close()
-	return f.r.Close()
+	return f.stream.Close()
 }

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
 }

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}
+}
+
+// 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
 }
 
-// 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)
+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
+}