#155 refactor and synchronise writing to CodeOcean. (#174)

* #155 refactor and synchronise writing to CodeOcean.

* Reduce complexity of input parsing.

* Update typo in internal/api/ws/codeocean_writer.go

Co-authored-by: Sebastian Serth <MrSerth@users.noreply.github.com>

internal/api/ws/codeocean_reader.go

@@ -0,0 +1,177 @@
+package ws
+
+import (
+	"context"
+	"github.com/gorilla/websocket"
+	"github.com/openHPI/poseidon/pkg/logging"
+	"io"
+)
+
+const CodeOceanToRawReaderBufferSize = 1024
+
+var log = logging.GetLogger("ws")
+
+// WebSocketReader is an interface that is intended for providing abstraction around reading from a WebSocket.
+// Besides, io.Reader, it also implements io.Writer. The Write method is used to inject data into the WebSocket stream.
+type WebSocketReader interface {
+	io.Reader
+	io.Writer
+	Start()
+	Stop()
+}
+
+// codeOceanToRawReader is an io.Reader implementation that provides the content of the WebSocket connection
+// to CodeOcean. You have to start the Reader by calling readInputLoop. After that you can use the Read function.
+type codeOceanToRawReader struct {
+	connection Connection
+
+	// readCtx is the context in that messages from CodeOcean are read.
+	readCtx       context.Context
+	cancelReadCtx context.CancelFunc
+	// executorCtx is the context in that messages are forwarded to the executor.
+	executorCtx context.Context
+
+	// A buffered channel of bytes is used to store data coming from CodeOcean via WebSocket
+	// and retrieve it when Read(..) is called. Since channels are thread-safe, we use one here
+	// instead of bytes.Buffer.
+	buffer chan byte
+	// The priorityBuffer is a buffer for injecting data into stdin of the execution from Poseidon,
+	// for example the character that causes the tty to generate a SIGQUIT signal.
+	// It is always read before the regular buffer.
+	priorityBuffer chan byte
+}
+
+func NewCodeOceanToRawReader(connection Connection, wsCtx, executorCtx context.Context) *codeOceanToRawReader {
+	return &codeOceanToRawReader{
+		connection:     connection,
+		readCtx:        wsCtx, // This context may be canceled before the executorCtx.
+		cancelReadCtx:  func() {},
+		executorCtx:    executorCtx,
+		buffer:         make(chan byte, CodeOceanToRawReaderBufferSize),
+		priorityBuffer: make(chan byte, CodeOceanToRawReaderBufferSize),
+	}
+}
+
+// readInputLoop reads from the WebSocket connection and buffers the user's input.
+// This is necessary because input must be read for the connection to handle special messages like close and call the
+// CloseHandler.
+func (cr *codeOceanToRawReader) readInputLoop(ctx context.Context) {
+	readMessage := make(chan bool)
+	loopContext, cancelInputLoop := context.WithCancel(ctx)
+	defer cancelInputLoop()
+	readingContext, cancelNextMessage := context.WithCancel(loopContext)
+	defer cancelNextMessage()
+
+	for loopContext.Err() == nil {
+		var messageType int
+		var reader io.Reader
+		var err error
+
+		go func() {
+			messageType, reader, err = cr.connection.NextReader()
+			select {
+			case <-readingContext.Done():
+			case readMessage <- true:
+			}
+		}()
+		select {
+		case <-loopContext.Done():
+			return
+		case <-readMessage:
+		}
+
+		if inputContainsError(messageType, err) {
+			return
+		}
+		if handleInput(reader, cr.buffer, loopContext) {
+			return
+		}
+	}
+}
+
+// handleInput receives a new message from the client and may forward it to the executor.
+func handleInput(reader io.Reader, buffer chan byte, ctx context.Context) (done bool) {
+	message, err := io.ReadAll(reader)
+	if err != nil {
+		log.WithError(err).Warn("error while reading WebSocket message")
+		return true
+	}
+
+	log.WithField("message", string(message)).Trace("Received message from client")
+	for _, character := range message {
+		select {
+		case <-ctx.Done():
+			return true
+		case buffer <- character:
+		}
+	}
+	return false
+}
+
+func inputContainsError(messageType int, err error) (done bool) {
+	if err != nil && websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+		log.Debug("ReadInputLoop: The client closed the connection!")
+		// The close handler will do something soon.
+		return true
+	} else if err != nil {
+		log.WithError(err).Warn("Error reading client message")
+		return true
+	}
+	if messageType != websocket.TextMessage {
+		log.WithField("messageType", messageType).Warn("Received message of wrong type")
+		return true
+	}
+	return false
+}
+
+// Start starts the read input loop asynchronously.
+func (cr *codeOceanToRawReader) Start() {
+	ctx, cancel := context.WithCancel(cr.readCtx)
+	cr.cancelReadCtx = cancel
+	go cr.readInputLoop(ctx)
+}
+
+// Stop stops the asynchronous read input loop.
+func (cr *codeOceanToRawReader) Stop() {
+	cr.cancelReadCtx()
+}
+
+// Read implements the io.Reader interface.
+// It returns bytes from the buffer or priorityBuffer.
+func (cr *codeOceanToRawReader) Read(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+
+	// Ensure to not return until at least one byte has been read to avoid busy waiting.
+	select {
+	case <-cr.executorCtx.Done():
+		return 0, io.EOF
+	case p[0] = <-cr.priorityBuffer:
+	case p[0] = <-cr.buffer:
+	}
+	var n int
+	for n = 1; n < len(p); n++ {
+		select {
+		case p[n] = <-cr.priorityBuffer:
+		case p[n] = <-cr.buffer:
+		default:
+			return n, nil
+		}
+	}
+	return n, nil
+}
+
+// Write implements the io.Writer interface.
+// Data written to a codeOceanToRawReader using this method is returned by Read before other data from the WebSocket.
+func (cr *codeOceanToRawReader) Write(p []byte) (n int, err error) {
+	var c byte
+	for n, c = range p {
+		select {
+		case cr.priorityBuffer <- c:
+		default:
+			break
+		}
+	}
+	return n, nil
+}

internal/api/ws/codeocean_reader_test.go

@@ -0,0 +1,75 @@
+package ws
+
+import (
+	"context"
+	"github.com/gorilla/websocket"
+	"github.com/openHPI/poseidon/tests"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"io"
+	"strings"
+	"testing"
+)
+
+func TestCodeOceanToRawReaderReturnsOnlyAfterOneByteWasRead(t *testing.T) {
+	readingCtx, cancel := context.WithCancel(context.Background())
+	forwardingCtx := readingCtx
+	defer cancel()
+	reader := NewCodeOceanToRawReader(nil, readingCtx, forwardingCtx)
+
+	read := make(chan bool)
+	go func() {
+		//nolint:makezero // we can't make zero initial length here as the reader otherwise doesn't block
+		p := make([]byte, 10)
+		_, err := reader.Read(p)
+		require.NoError(t, err)
+		read <- true
+	}()
+
+	t.Run("Does not return immediately when there is no data", func(t *testing.T) {
+		assert.False(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout))
+	})
+
+	t.Run("Returns when there is data available", func(t *testing.T) {
+		reader.buffer <- byte(42)
+		assert.True(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout))
+	})
+}
+
+func TestCodeOceanToRawReaderReturnsOnlyAfterOneByteWasReadFromConnection(t *testing.T) {
+	messages := make(chan io.Reader)
+
+	connection := &ConnectionMock{}
+	connection.On("WriteMessage", mock.AnythingOfType("int"), mock.AnythingOfType("[]uint8")).Return(nil)
+	connection.On("CloseHandler").Return(nil)
+	connection.On("SetCloseHandler", mock.Anything).Return()
+	call := connection.On("NextReader")
+	call.Run(func(_ mock.Arguments) {
+		call.Return(websocket.TextMessage, <-messages, nil)
+	})
+
+	readingCtx, cancel := context.WithCancel(context.Background())
+	forwardingCtx := readingCtx
+	defer cancel()
+	reader := NewCodeOceanToRawReader(connection, readingCtx, forwardingCtx)
+	reader.Start()
+
+	read := make(chan bool)
+	//nolint:makezero // this is required here to make the Read call blocking
+	message := make([]byte, 10)
+	go func() {
+		_, err := reader.Read(message)
+		require.NoError(t, err)
+		read <- true
+	}()
+
+	t.Run("Does not return immediately when there is no data", func(t *testing.T) {
+		assert.False(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout))
+	})
+
+	t.Run("Returns when there is data available", func(t *testing.T) {
+		messages <- strings.NewReader("Hello")
+		assert.True(t, tests.ChannelReceivesSomething(read, tests.ShortTimeout))
+	})
+}

internal/api/ws/codeocean_writer.go

@@ -0,0 +1,155 @@
+package ws
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"github.com/gorilla/websocket"
+	"github.com/openHPI/poseidon/internal/runner"
+	"github.com/openHPI/poseidon/pkg/dto"
+	"io"
+)
+
+// CodeOceanOutputWriterBufferSize defines the number of messages.
+const CodeOceanOutputWriterBufferSize = 64
+
+// rawToCodeOceanWriter is a simple io.Writer implementation that just forwards the call to sendMessage.
+type rawToCodeOceanWriter struct {
+	sendMessage func(string)
+}
+
+// Write implements the io.Writer interface.
+func (rc *rawToCodeOceanWriter) Write(p []byte) (int, error) {
+	rc.sendMessage(string(p))
+	return len(p), nil
+}
+
+// WebSocketWriter is an interface that defines which data is required and which information can be passed.
+type WebSocketWriter interface {
+	StdOut() io.Writer
+	StdErr() io.Writer
+	SendExitInfo(info *runner.ExitInfo)
+}
+
+// codeOceanOutputWriter is a concrete WebSocketWriter implementation.
+// It forwards the data written to stdOut or stdErr (Nomad, AWS) to the WebSocket connection (CodeOcean).
+type codeOceanOutputWriter struct {
+	connection Connection
+	stdOut     io.Writer
+	stdErr     io.Writer
+	queue      chan *writingLoopMessage
+	stopped    bool
+}
+
+// writingLoopMessage is an internal data structure to notify the writing loop when it should stop.
+type writingLoopMessage struct {
+	done bool
+	data *dto.WebSocketMessage
+}
+
+// NewCodeOceanOutputWriter provies an codeOceanOutputWriter for the time the context ctx is active.
+// The codeOceanOutputWriter handles all the messages defined in the websocket.schema.json (start, timeout, stdout, ..).
+func NewCodeOceanOutputWriter(connection Connection, ctx context.Context) *codeOceanOutputWriter {
+	cw := &codeOceanOutputWriter{
+		connection: connection,
+		queue:      make(chan *writingLoopMessage, CodeOceanOutputWriterBufferSize),
+		stopped:    false,
+	}
+	cw.stdOut = &rawToCodeOceanWriter{sendMessage: func(s string) {
+		cw.send(&dto.WebSocketMessage{Type: dto.WebSocketOutputStdout, Data: s})
+	}}
+	cw.stdErr = &rawToCodeOceanWriter{sendMessage: func(s string) {
+		cw.send(&dto.WebSocketMessage{Type: dto.WebSocketOutputStderr, Data: s})
+	}}
+
+	go cw.startWritingLoop()
+	go cw.stopWhenContextDone(ctx)
+	cw.send(&dto.WebSocketMessage{Type: dto.WebSocketMetaStart})
+	return cw
+}
+
+// StdOut provides an io.Writer that forwards the written data to CodeOcean as StdOut stream.
+func (cw *codeOceanOutputWriter) StdOut() io.Writer {
+	return cw.stdOut
+}
+
+// StdErr provides an io.Writer that forwards the written data to CodeOcean as StdErr stream.
+func (cw *codeOceanOutputWriter) StdErr() io.Writer {
+	return cw.stdErr
+}
+
+// SendExitInfo forwards the kind of exit (timeout, error, normal) to CodeOcean.
+func (cw *codeOceanOutputWriter) SendExitInfo(info *runner.ExitInfo) {
+	switch {
+	case errors.Is(info.Err, context.DeadlineExceeded) || errors.Is(info.Err, runner.ErrorRunnerInactivityTimeout):
+		cw.send(&dto.WebSocketMessage{Type: dto.WebSocketMetaTimeout})
+	case info.Err != nil:
+		errorMessage := "Error executing the request"
+		log.WithError(info.Err).Warn(errorMessage)
+		cw.send(&dto.WebSocketMessage{Type: dto.WebSocketOutputError, Data: errorMessage})
+	default:
+		cw.send(&dto.WebSocketMessage{Type: dto.WebSocketExit, ExitCode: info.Code})
+	}
+}
+
+// stopWhenContextDone notifies the writing loop to stop after the context has been passed.
+func (cw *codeOceanOutputWriter) stopWhenContextDone(ctx context.Context) {
+	<-ctx.Done()
+	if !cw.stopped {
+		cw.queue <- &writingLoopMessage{done: true}
+	}
+}
+
+// send forwards the passed dto.WebSocketMessage to the writing loop.
+func (cw *codeOceanOutputWriter) send(message *dto.WebSocketMessage) {
+	if cw.stopped {
+		return
+	}
+	done := message.Type == dto.WebSocketExit ||
+		message.Type == dto.WebSocketMetaTimeout ||
+		message.Type == dto.WebSocketOutputError
+	cw.queue <- &writingLoopMessage{done: done, data: message}
+}
+
+// startWritingLoop enables the writing loop.
+// This is the central and only place where written changes to the WebSocket connection should be done.
+// It synchronizes the messages to provide state checks of the WebSocket connection.
+func (cw *codeOceanOutputWriter) startWritingLoop() {
+	for {
+		message := <-cw.queue
+		done := sendMessage(cw.connection, message.data)
+		if done || message.done {
+			break
+		}
+	}
+	cw.stopped = true
+	message := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
+	err := cw.connection.WriteMessage(websocket.CloseMessage, message)
+	err2 := cw.connection.Close()
+	if err != nil || err2 != nil {
+		log.WithError(err).WithField("err2", err2).Warn("Error during websocket close")
+	}
+}
+
+// sendMessage is a helper function for the writing loop. It must not be called from somewhere else!
+func sendMessage(connection Connection, message *dto.WebSocketMessage) (done bool) {
+	if message == nil {
+		return false
+	}
+
+	encodedMessage, err := json.Marshal(message)
+	if err != nil {
+		log.WithField("message", message).WithError(err).Warn("Marshal error")
+		return false
+	}
+
+	log.WithField("message", message).Trace("Sending message to client")
+	err = connection.WriteMessage(websocket.TextMessage, encodedMessage)
+	if err != nil {
+		errorMessage := "Error writing the message"
+		log.WithField("message", message).WithError(err).Warn(errorMessage)
+		return true
+	}
+
+	return false
+}

internal/api/ws/codeocean_writer_test.go

@@ -0,0 +1,100 @@
+package ws
+
+import (
+	"context"
+	"encoding/json"
+	"github.com/gorilla/websocket"
+	"github.com/openHPI/poseidon/internal/runner"
+	"github.com/openHPI/poseidon/pkg/dto"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestRawToCodeOceanWriter(t *testing.T) {
+	connectionMock, message := buildConnectionMock(t)
+	proxyCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	output := NewCodeOceanOutputWriter(connectionMock, proxyCtx)
+	<-message // start message
+
+	t.Run("StdOut", func(t *testing.T) {
+		testMessage := "testStdOut"
+		_, err := output.StdOut().Write([]byte(testMessage))
+		require.NoError(t, err)
+
+		expected, err := json.Marshal(struct {
+			Type string `json:"type"`
+			Data string `json:"data"`
+		}{string(dto.WebSocketOutputStdout), testMessage})
+		require.NoError(t, err)
+
+		assert.Equal(t, expected, <-message)
+	})
+
+	t.Run("StdErr", func(t *testing.T) {
+		testMessage := "testStdErr"
+		_, err := output.StdErr().Write([]byte(testMessage))
+		require.NoError(t, err)
+
+		expected, err := json.Marshal(struct {
+			Type string `json:"type"`
+			Data string `json:"data"`
+		}{string(dto.WebSocketOutputStderr), testMessage})
+		require.NoError(t, err)
+
+		assert.Equal(t, expected, <-message)
+	})
+}
+
+type sendExitInfoTestCase struct {
+	name    string
+	info    *runner.ExitInfo
+	message dto.WebSocketMessage
+}
+
+func TestCodeOceanOutputWriter_SendExitInfo(t *testing.T) {
+	testCases := []sendExitInfoTestCase{
+		{"Timeout", &runner.ExitInfo{Err: runner.ErrorRunnerInactivityTimeout},
+			dto.WebSocketMessage{Type: dto.WebSocketMetaTimeout}},
+		{"Error", &runner.ExitInfo{Err: websocket.ErrCloseSent},
+			dto.WebSocketMessage{Type: dto.WebSocketOutputError, Data: "Error executing the request"}},
+		{"Exit", &runner.ExitInfo{Code: 21},
+			dto.WebSocketMessage{Type: dto.WebSocketExit, ExitCode: 21}},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.name, func(t *testing.T) {
+			connectionMock, message := buildConnectionMock(t)
+			proxyCtx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			output := NewCodeOceanOutputWriter(connectionMock, proxyCtx)
+			<-message // start message
+
+			output.SendExitInfo(test.info)
+			expected, err := json.Marshal(test.message)
+			require.NoError(t, err)
+
+			msg := <-message
+			assert.Equal(t, expected, msg)
+		})
+	}
+}
+
+func buildConnectionMock(t *testing.T) (conn *ConnectionMock, messages chan []byte) {
+	t.Helper()
+	message := make(chan []byte)
+	connectionMock := &ConnectionMock{}
+	connectionMock.On("WriteMessage", mock.AnythingOfType("int"), mock.AnythingOfType("[]uint8")).
+		Run(func(args mock.Arguments) {
+			m, ok := args.Get(1).([]byte)
+			require.True(t, ok)
+			message <- m
+		}).
+		Return(nil)
+	connectionMock.On("CloseHandler").Return(nil)
+	connectionMock.On("SetCloseHandler", mock.Anything).Return()
+	connectionMock.On("Close").Return()
+	return connectionMock, message
+}

internal/api/ws/connection.go

@@ -0,0 +1,14 @@
+package ws
+
+import (
+	"io"
+)
+
+// Connection is an internal interface for websocket.Conn in order to mock it for unit tests.
+type Connection interface {
+	WriteMessage(messageType int, data []byte) error
+	Close() error
+	NextReader() (messageType int, r io.Reader, err error)
+	CloseHandler() func(code int, text string) error
+	SetCloseHandler(handler func(code int, text string) error)
+}

internal/api/ws/connection_mock.go

@@ -0,0 +1,108 @@
+// Code generated by mockery v2.13.1. DO NOT EDIT.
+
+package ws
+
+import (
+	io "io"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// ConnectionMock is an autogenerated mock type for the Connection type
+type ConnectionMock struct {
+	mock.Mock
+}
+
+// Close provides a mock function with given fields:
+func (_m *ConnectionMock) Close() error {
+	ret := _m.Called()
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// CloseHandler provides a mock function with given fields:
+func (_m *ConnectionMock) CloseHandler() func(int, string) error {
+	ret := _m.Called()
+
+	var r0 func(int, string) error
+	if rf, ok := ret.Get(0).(func() func(int, string) error); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(func(int, string) error)
+		}
+	}
+
+	return r0
+}
+
+// NextReader provides a mock function with given fields:
+func (_m *ConnectionMock) NextReader() (int, io.Reader, error) {
+	ret := _m.Called()
+
+	var r0 int
+	if rf, ok := ret.Get(0).(func() int); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(int)
+	}
+
+	var r1 io.Reader
+	if rf, ok := ret.Get(1).(func() io.Reader); ok {
+		r1 = rf()
+	} else {
+		if ret.Get(1) != nil {
+			r1 = ret.Get(1).(io.Reader)
+		}
+	}
+
+	var r2 error
+	if rf, ok := ret.Get(2).(func() error); ok {
+		r2 = rf()
+	} else {
+		r2 = ret.Error(2)
+	}
+
+	return r0, r1, r2
+}
+
+// SetCloseHandler provides a mock function with given fields: handler
+func (_m *ConnectionMock) SetCloseHandler(handler func(int, string) error) {
+	_m.Called(handler)
+}
+
+// WriteMessage provides a mock function with given fields: messageType, data
+func (_m *ConnectionMock) WriteMessage(messageType int, data []byte) error {
+	ret := _m.Called(messageType, data)
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(int, []byte) error); ok {
+		r0 = rf(messageType, data)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+type mockConstructorTestingTNewConnectionMock interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewConnectionMock creates a new instance of ConnectionMock. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewConnectionMock(t mockConstructorTestingTNewConnectionMock) *ConnectionMock {
+	mock := &ConnectionMock{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}