2e701340e5
- Add protocol buffer optimizations (internal/api/protocol.go) - Add filesystem queue backend (internal/queue/filesystem_queue.go) - Add run manifest support (internal/manifest/run_manifest.go) - Worker and jupyter task refinements - Exported test wrappers for benchmarking
358 lines
8.8 KiB
Go
358 lines
8.8 KiB
Go
package api
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"encoding/json"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var bufferPool = sync.Pool{
|
|
New: func() interface{} {
|
|
buf := make([]byte, 0, 256)
|
|
return &buf
|
|
},
|
|
}
|
|
|
|
// Response packet types
|
|
const (
|
|
PacketTypeSuccess = 0x00
|
|
PacketTypeError = 0x01
|
|
PacketTypeProgress = 0x02
|
|
PacketTypeStatus = 0x03
|
|
PacketTypeData = 0x04
|
|
PacketTypeLog = 0x05
|
|
)
|
|
|
|
// Error codes
|
|
const (
|
|
ErrorCodeUnknownError = 0x00
|
|
ErrorCodeInvalidRequest = 0x01
|
|
ErrorCodeAuthenticationFailed = 0x02
|
|
ErrorCodePermissionDenied = 0x03
|
|
ErrorCodeResourceNotFound = 0x04
|
|
ErrorCodeResourceAlreadyExists = 0x05
|
|
|
|
ErrorCodeServerOverloaded = 0x10
|
|
ErrorCodeDatabaseError = 0x11
|
|
ErrorCodeNetworkError = 0x12
|
|
ErrorCodeStorageError = 0x13
|
|
ErrorCodeTimeout = 0x14
|
|
|
|
ErrorCodeJobNotFound = 0x20
|
|
ErrorCodeJobAlreadyRunning = 0x21
|
|
ErrorCodeJobFailedToStart = 0x22
|
|
ErrorCodeJobExecutionFailed = 0x23
|
|
ErrorCodeJobCancelled = 0x24
|
|
|
|
ErrorCodeOutOfMemory = 0x30
|
|
ErrorCodeDiskFull = 0x31
|
|
ErrorCodeInvalidConfiguration = 0x32
|
|
ErrorCodeServiceUnavailable = 0x33
|
|
)
|
|
|
|
// Progress types
|
|
const (
|
|
ProgressTypePercentage = 0x00
|
|
ProgressTypeStage = 0x01
|
|
ProgressTypeMessage = 0x02
|
|
ProgressTypeBytesTransferred = 0x03
|
|
)
|
|
|
|
// Log levels
|
|
const (
|
|
LogLevelDebug = 0x00
|
|
LogLevelInfo = 0x01
|
|
LogLevelWarn = 0x02
|
|
LogLevelError = 0x03
|
|
)
|
|
|
|
// ResponsePacket represents a structured response packet
|
|
type ResponsePacket struct {
|
|
PacketType byte
|
|
Timestamp uint64
|
|
|
|
// Success fields
|
|
SuccessMessage string
|
|
|
|
// Error fields
|
|
ErrorCode byte
|
|
ErrorMessage string
|
|
ErrorDetails string
|
|
|
|
// Progress fields
|
|
ProgressType byte
|
|
ProgressValue uint32
|
|
ProgressTotal uint32
|
|
ProgressMessage string
|
|
|
|
// Status fields
|
|
StatusData string
|
|
|
|
// Data fields
|
|
DataType string
|
|
DataPayload []byte
|
|
|
|
// Log fields
|
|
LogLevel byte
|
|
LogMessage string
|
|
}
|
|
|
|
// NewSuccessPacket creates a success response packet
|
|
func NewSuccessPacket(message string) *ResponsePacket {
|
|
return &ResponsePacket{
|
|
PacketType: PacketTypeSuccess,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
SuccessMessage: message,
|
|
}
|
|
}
|
|
|
|
// NewSuccessPacketWithPayload creates a success response packet with JSON payload
|
|
func NewSuccessPacketWithPayload(message string, payload interface{}) *ResponsePacket {
|
|
// Convert payload to JSON for the DataPayload field
|
|
payloadBytes, _ := json.Marshal(payload)
|
|
|
|
return &ResponsePacket{
|
|
PacketType: PacketTypeData,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
SuccessMessage: message,
|
|
DataType: "status",
|
|
DataPayload: payloadBytes,
|
|
}
|
|
}
|
|
|
|
// NewErrorPacket creates an error response packet
|
|
func NewErrorPacket(errorCode byte, message string, details string) *ResponsePacket {
|
|
return &ResponsePacket{
|
|
PacketType: PacketTypeError,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ErrorCode: errorCode,
|
|
ErrorMessage: message,
|
|
ErrorDetails: details,
|
|
}
|
|
}
|
|
|
|
// NewProgressPacket creates a progress response packet
|
|
func NewProgressPacket(
|
|
progressType byte,
|
|
value uint32,
|
|
total uint32,
|
|
message string,
|
|
) *ResponsePacket {
|
|
return &ResponsePacket{
|
|
PacketType: PacketTypeProgress,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
ProgressType: progressType,
|
|
ProgressValue: value,
|
|
ProgressTotal: total,
|
|
ProgressMessage: message,
|
|
}
|
|
}
|
|
|
|
// NewStatusPacket creates a status response packet
|
|
func NewStatusPacket(data string) *ResponsePacket {
|
|
return &ResponsePacket{
|
|
PacketType: PacketTypeStatus,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
StatusData: data,
|
|
}
|
|
}
|
|
|
|
// NewDataPacket creates a data response packet
|
|
func NewDataPacket(dataType string, payload []byte) *ResponsePacket {
|
|
return &ResponsePacket{
|
|
PacketType: PacketTypeData,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
DataType: dataType,
|
|
DataPayload: payload,
|
|
}
|
|
}
|
|
|
|
// NewLogPacket creates a log response packet
|
|
func NewLogPacket(level byte, message string) *ResponsePacket {
|
|
return &ResponsePacket{
|
|
PacketType: PacketTypeLog,
|
|
Timestamp: uint64(time.Now().Unix()),
|
|
LogLevel: level,
|
|
LogMessage: message,
|
|
}
|
|
}
|
|
|
|
// Serialize converts the packet to binary format
|
|
func (p *ResponsePacket) Serialize() ([]byte, error) {
|
|
// For small packets, avoid pool overhead
|
|
if p.estimatedSize() <= 1024 {
|
|
buf := make([]byte, 0, p.estimatedSize())
|
|
return serializePacketToBuffer(p, buf)
|
|
}
|
|
|
|
// Use pool for larger packets
|
|
bufPtr := bufferPool.Get().(*[]byte)
|
|
defer func() {
|
|
*bufPtr = (*bufPtr)[:0]
|
|
bufferPool.Put(bufPtr)
|
|
}()
|
|
buf := *bufPtr
|
|
|
|
// Ensure buffer has enough capacity
|
|
if cap(buf) < p.estimatedSize() {
|
|
buf = make([]byte, 0, p.estimatedSize())
|
|
} else {
|
|
buf = buf[:0]
|
|
}
|
|
|
|
return serializePacketToBuffer(p, buf)
|
|
}
|
|
|
|
func serializePacketToBuffer(p *ResponsePacket, buf []byte) ([]byte, error) {
|
|
// Packet type
|
|
buf = append(buf, p.PacketType)
|
|
|
|
// Timestamp (8 bytes, big-endian)
|
|
var timestampBytes [8]byte
|
|
binary.BigEndian.PutUint64(timestampBytes[:], p.Timestamp)
|
|
buf = append(buf, timestampBytes[:]...)
|
|
|
|
// Packet-specific data
|
|
switch p.PacketType {
|
|
case PacketTypeSuccess:
|
|
buf = appendString(buf, p.SuccessMessage)
|
|
|
|
case PacketTypeError:
|
|
buf = append(buf, p.ErrorCode)
|
|
buf = appendString(buf, p.ErrorMessage)
|
|
buf = appendString(buf, p.ErrorDetails)
|
|
|
|
case PacketTypeProgress:
|
|
buf = append(buf, p.ProgressType)
|
|
buf = appendUint32(buf, p.ProgressValue)
|
|
buf = appendUint32(buf, p.ProgressTotal)
|
|
buf = appendString(buf, p.ProgressMessage)
|
|
|
|
case PacketTypeStatus:
|
|
buf = appendString(buf, p.StatusData)
|
|
|
|
case PacketTypeData:
|
|
buf = appendString(buf, p.DataType)
|
|
buf = appendBytes(buf, p.DataPayload)
|
|
|
|
case PacketTypeLog:
|
|
buf = append(buf, p.LogLevel)
|
|
buf = appendString(buf, p.LogMessage)
|
|
|
|
default:
|
|
return nil, fmt.Errorf("unknown packet type: %d", p.PacketType)
|
|
}
|
|
|
|
return buf, nil
|
|
}
|
|
|
|
// appendString writes a string with varint length prefix
|
|
func appendString(buf []byte, s string) []byte {
|
|
length := uint64(len(s))
|
|
var tmp [binary.MaxVarintLen64]byte
|
|
n := binary.PutUvarint(tmp[:], length)
|
|
buf = append(buf, tmp[:n]...)
|
|
return append(buf, s...)
|
|
}
|
|
|
|
// appendBytes writes bytes with varint length prefix
|
|
func appendBytes(buf []byte, b []byte) []byte {
|
|
length := uint64(len(b))
|
|
var tmp [binary.MaxVarintLen64]byte
|
|
n := binary.PutUvarint(tmp[:], length)
|
|
buf = append(buf, tmp[:n]...)
|
|
return append(buf, b...)
|
|
}
|
|
|
|
func appendUint32(buf []byte, value uint32) []byte {
|
|
var tmp [4]byte
|
|
binary.BigEndian.PutUint32(tmp[:], value)
|
|
return append(buf, tmp[:]...)
|
|
}
|
|
|
|
func (p *ResponsePacket) estimatedSize() int {
|
|
base := 1 + 8 // packet type + timestamp
|
|
|
|
switch p.PacketType {
|
|
case PacketTypeSuccess:
|
|
return base + binary.MaxVarintLen64 + len(p.SuccessMessage)
|
|
case PacketTypeError:
|
|
return base + 1 + 2*binary.MaxVarintLen64 + len(p.ErrorMessage) + len(p.ErrorDetails)
|
|
case PacketTypeProgress:
|
|
return base + 1 + 4 + 4 + binary.MaxVarintLen64 + len(p.ProgressMessage)
|
|
case PacketTypeStatus:
|
|
return base + binary.MaxVarintLen64 + len(p.StatusData)
|
|
case PacketTypeData:
|
|
return base + binary.MaxVarintLen64 + len(p.DataType) + binary.MaxVarintLen64 + len(p.DataPayload)
|
|
case PacketTypeLog:
|
|
return base + 1 + binary.MaxVarintLen64 + len(p.LogMessage)
|
|
default:
|
|
return base
|
|
}
|
|
}
|
|
|
|
// GetErrorMessage returns a human-readable error message for an error code
|
|
func GetErrorMessage(code byte) string {
|
|
switch code {
|
|
case ErrorCodeUnknownError:
|
|
return "Unknown error occurred"
|
|
case ErrorCodeInvalidRequest:
|
|
return "Invalid request format"
|
|
case ErrorCodeAuthenticationFailed:
|
|
return "Authentication failed"
|
|
case ErrorCodePermissionDenied:
|
|
return "Permission denied"
|
|
case ErrorCodeResourceNotFound:
|
|
return "Resource not found"
|
|
case ErrorCodeResourceAlreadyExists:
|
|
return "Resource already exists"
|
|
case ErrorCodeServerOverloaded:
|
|
return "Server is overloaded"
|
|
case ErrorCodeDatabaseError:
|
|
return "Database error occurred"
|
|
case ErrorCodeNetworkError:
|
|
return "Network error occurred"
|
|
case ErrorCodeStorageError:
|
|
return "Storage error occurred"
|
|
case ErrorCodeTimeout:
|
|
return "Operation timed out"
|
|
case ErrorCodeJobNotFound:
|
|
return "Job not found"
|
|
case ErrorCodeJobAlreadyRunning:
|
|
return "Job is already running"
|
|
case ErrorCodeJobFailedToStart:
|
|
return "Job failed to start"
|
|
case ErrorCodeJobExecutionFailed:
|
|
return "Job execution failed"
|
|
case ErrorCodeJobCancelled:
|
|
return "Job was cancelled"
|
|
case ErrorCodeOutOfMemory:
|
|
return "Server out of memory"
|
|
case ErrorCodeDiskFull:
|
|
return "Server disk full"
|
|
case ErrorCodeInvalidConfiguration:
|
|
return "Invalid server configuration"
|
|
case ErrorCodeServiceUnavailable:
|
|
return "Service temporarily unavailable"
|
|
default:
|
|
return "Unknown error code"
|
|
}
|
|
}
|
|
|
|
// GetLogLevelName returns the name for a log level
|
|
func GetLogLevelName(level byte) string {
|
|
switch level {
|
|
case LogLevelDebug:
|
|
return "DEBUG"
|
|
case LogLevelInfo:
|
|
return "INFO"
|
|
case LogLevelWarn:
|
|
return "WARN"
|
|
case LogLevelError:
|
|
return "ERROR"
|
|
default:
|
|
return "UNKNOWN"
|
|
}
|
|
}
|