added poseidon with aws to k8s changes

This commit is contained in:
Elmar Kresse
2024-08-12 10:02:36 +02:00
parent 5376f7a027
commit 254460d64c
60 changed files with 6912 additions and 0 deletions

View File

@@ -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, loopContext) {
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.WithContext(ctx).WithError(err).Warn("error while reading WebSocket message")
return true
}
log.WithContext(ctx).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, ctx context.Context) (done bool) {
if err != nil && websocket.IsCloseError(err, websocket.CloseNormalClosure) {
log.WithContext(ctx).Debug("ReadInputLoop: The client closed the connection!")
// The close handler will do something soon.
return true
} else if err != nil {
log.WithContext(ctx).WithError(err).Warn("Error reading client message")
return true
}
if messageType != websocket.TextMessage {
log.WithContext(ctx).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
}

View File

@@ -0,0 +1,174 @@
package ws
import (
"context"
"encoding/json"
"errors"
"fmt"
"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 {
outputType dto.WebSocketMessageType
sendMessage func(*dto.WebSocketMessage)
ctx context.Context
}
// Write implements the io.Writer interface.
func (rc *rawToCodeOceanWriter) Write(p []byte) (int, error) {
switch {
case rc.ctx.Err() != nil:
return 0, fmt.Errorf("CodeOceanWriter context done: %w", rc.ctx.Err())
case len(p) == 0:
return 0, nil
default:
rc.sendMessage(&dto.WebSocketMessage{Type: rc.outputType, Data: 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
Close(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 *rawToCodeOceanWriter
stdErr *rawToCodeOceanWriter
queue chan *writingLoopMessage
ctx context.Context
}
// writingLoopMessage is an internal data structure to notify the writing loop when it should stop.
type writingLoopMessage struct {
done bool
data *dto.WebSocketMessage
}
// NewCodeOceanOutputWriter provides 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, done context.CancelFunc) WebSocketWriter {
cw := &codeOceanOutputWriter{
connection: connection,
queue: make(chan *writingLoopMessage, CodeOceanOutputWriterBufferSize),
ctx: ctx,
}
cw.stdOut = &rawToCodeOceanWriter{
outputType: dto.WebSocketOutputStdout,
sendMessage: cw.send,
ctx: ctx,
}
cw.stdErr = &rawToCodeOceanWriter{
outputType: dto.WebSocketOutputStderr,
sendMessage: cw.send,
ctx: ctx,
}
go cw.startWritingLoop(done)
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
}
// Close forwards the kind of exit (timeout, error, normal) to CodeOcean.
// This results in the closing of the WebSocket connection.
// The call of Close is mandatory!
func (cw *codeOceanOutputWriter) Close(info *runner.ExitInfo) {
defer func() { cw.queue <- &writingLoopMessage{done: true} }()
// Mask the internal stop reason before disclosing/forwarding it externally/to CodeOcean.
switch {
case info == nil:
return
case info.Err == nil:
cw.send(&dto.WebSocketMessage{Type: dto.WebSocketExit, ExitCode: info.Code})
case errors.Is(info.Err, runner.ErrorExecutionTimeout) || errors.Is(info.Err, runner.ErrorRunnerInactivityTimeout):
cw.send(&dto.WebSocketMessage{Type: dto.WebSocketMetaTimeout})
case errors.Is(info.Err, runner.ErrOOMKilled):
cw.send(&dto.WebSocketMessage{Type: dto.WebSocketOutputError, Data: dto.ErrOOMKilled.Error()})
case errors.Is(info.Err, runner.ErrDestroyedByAPIRequest):
message := "the allocation stopped as expected"
log.WithContext(cw.ctx).WithError(info.Err).Trace(message)
cw.send(&dto.WebSocketMessage{Type: dto.WebSocketOutputError, Data: message})
default:
errorMessage := "Error executing the request"
log.WithContext(cw.ctx).WithError(info.Err).Warn(errorMessage)
cw.send(&dto.WebSocketMessage{Type: dto.WebSocketOutputError, Data: errorMessage})
}
}
// send forwards the passed dto.WebSocketMessage to the writing loop.
func (cw *codeOceanOutputWriter) send(message *dto.WebSocketMessage) {
cw.queue <- &writingLoopMessage{done: false, 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(writingLoopDone context.CancelFunc) {
defer func() {
message := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
err := cw.connection.WriteMessage(websocket.CloseMessage, message)
err2 := cw.connection.Close()
if err != nil || err2 != nil {
log.WithContext(cw.ctx).WithError(err).WithField("err2", err2).Warn("Error during websocket close")
}
}()
for {
message := <-cw.queue
done := true
if message.data != nil {
done = sendMessage(cw.connection, message.data, cw.ctx)
}
if done || message.done {
log.WithContext(cw.ctx).Trace("Writing loop done")
writingLoopDone()
return
}
}
}
// sendMessage is a helper function for the writing loop. It must not be called from somewhere else!
func sendMessage(connection Connection, message *dto.WebSocketMessage, ctx context.Context) (done bool) {
if message == nil {
return false
}
encodedMessage, err := json.Marshal(message)
if err != nil {
log.WithContext(ctx).WithField("message", message).WithError(err).Warn("Marshal error")
return false
}
log.WithContext(ctx).WithField("message", message).Trace("Sending message to client")
err = connection.WriteMessage(websocket.TextMessage, encodedMessage)
if err != nil {
errorMessage := "Error writing the message"
log.WithContext(ctx).WithField("message", message).WithError(err).Warn(errorMessage)
return true
}
return false
}

View File

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