feat(framing): add frame length to support large data (#191)

* feat(framing): add frame length and payload to support large data

* fix(client): auto re-connect when the connection is interrupted

Author: xiaojian-hong

pkg/client/client.go

@@ -11,6 +11,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/yomorun/yomo/pkg/framing"
 	"github.com/yomorun/yomo/pkg/quic"
 )
 
@@ -135,7 +136,6 @@ func (c *clientImpl) handleSignal(accepted chan bool) {
 				break
 			}
 			value := buf[:n]
-
 			if bytes.Equal(value, quic.SignalHeartbeat) {
 				// heartbeart
 				c.conn.Heartbeat <- buf[0]
@@ -169,9 +169,11 @@ func (c *clientImpl) handleSignal(accepted chan bool) {
 }
 
 // Write the data to downstream.
-func (c *clientImpl) Write(b []byte) (int, error) {
+func (c *clientImpl) Write(data []byte) (int, error) {
 	if c.conn.Stream != nil {
-		return c.conn.Stream.Write(b)
+		// wrap data with framing.
+		f := framing.NewPayloadFrame(data)
+		return c.conn.Stream.Write(f.Bytes())
 	} else {
 		return 0, errors.New("not found stream")
 	}

pkg/client/serverless.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"log"
 
+	"github.com/yomorun/yomo/pkg/framing"
 	"github.com/yomorun/yomo/pkg/quic"
 	"github.com/yomorun/yomo/pkg/rx"
 )
@@ -60,7 +61,10 @@ func (c *serverlessClientImpl) Pipe(f func(rxstream rx.RxStream) rx.RxStream) {
 				log.Print("❌ Please add the encode/marshal operator in the end of your Serverless handler.")
 				continue
 			}
-			_, err := c.writer.Write(buf)
+
+			// wrap data with framing.
+			f := framing.NewPayloadFrame(buf)
+			_, err := c.writer.Write(f.Bytes())
 			if err != nil {
 				log.Print("❌ Send data to zipper failed. ", err)
 			}

pkg/decoder/decoder.go

@@ -123,15 +123,14 @@ func FromStream(reader io.Reader) Observable {
 
 	f := func(next chan interface{}) {
 		defer close(next)
+		fd := NewFrameDecoder(reader)
 		for {
-			buf := make([]byte, 3*1024)
-			n, err := reader.Read(buf)
-
+			// read next raw frame.
+			buf, err := fd.Read(true)
 			if err != nil {
 				break
 			} else {
-				value := buf[:n]
-				next <- value
+				next <- buf
 			}
 		}
 	}

pkg/decoder/frame.go

@@ -0,0 +1,67 @@
+package decoder
+
+import (
+	"bufio"
+	"errors"
+	"io"
+
+	"github.com/yomorun/yomo/pkg/framing"
+)
+
+const (
+	minBuffSize = 3 * 1024
+	maxBuffSize = 16*1024*1024 + framing.FrameLengthFieldSize
+)
+
+// FrameDecoder defines a decoder for decoding frames which have a header of length.
+type FrameDecoder bufio.Scanner
+
+// Read reads next frame in bytes.
+// if rawFrame == true, it returns the raw frame without frame length field.
+func (p *FrameDecoder) Read(rawFrame bool) ([]byte, error) {
+	scanner := (*bufio.Scanner)(p)
+	if !scanner.Scan() {
+		err := scanner.Err()
+		if err == nil {
+			err = io.EOF
+		}
+		return nil, err
+	}
+
+	// read raw frame
+	if rawFrame {
+		raw := scanner.Bytes()[framing.FrameLengthFieldSize:]
+		return raw, nil
+	}
+
+	// read full frame includes frame length field.
+	return scanner.Bytes(), nil
+}
+
+func doSplit(data []byte, eof bool) (advance int, token []byte, err error) {
+	if eof {
+		return
+	}
+	if len(data) < framing.FrameLengthFieldSize {
+		return
+	}
+	frameLength := framing.ReadFrameLength(data)
+	if frameLength < 1 {
+		err = errors.New("invalid frame length")
+		return
+	}
+	frameSize := frameLength + framing.FrameLengthFieldSize
+	if frameSize <= len(data) {
+		return frameSize, data[:frameSize], nil
+	}
+	return
+}
+
+// NewFrameDecoder creates a new frame decoder.
+func NewFrameDecoder(r io.Reader) *FrameDecoder {
+	scanner := bufio.NewScanner(r)
+	scanner.Split(doSplit)
+	buf := make([]byte, 0, minBuffSize)
+	scanner.Buffer(buf, maxBuffSize)
+	return (*FrameDecoder)(scanner)
+}

pkg/framing/frame.go

@@ -0,0 +1,47 @@
+package framing
+
+// Frame represents a YoMo frame.
+//  +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+//  |                    Frame Length               |
+//  +-----------------------------------------------+
+//  |                     Frame                     |
+//  |                                               |
+//  +-----------------------------------------------+
+type Frame interface {
+	// Bytes converts the frame to bytes
+	Bytes() []byte
+}
+
+const FrameLengthFieldSize = 3
+
+func appendFrameLength(buf []byte, len int) {
+	for i := 0; i < FrameLengthFieldSize; i++ {
+		offset := 8 * (FrameLengthFieldSize - i - 1)
+		if offset > 0 {
+			buf[i] = byte(len >> offset)
+		} else {
+			buf[i] = byte(len)
+		}
+	}
+}
+
+// ReadFrameLength reads frame length from bytes.
+func ReadFrameLength(buf []byte) int {
+	c := 0
+	for i := 0; i < FrameLengthFieldSize; i++ {
+		offset := 8 * (FrameLengthFieldSize - i - 1)
+		if offset > 0 {
+			c += int(buf[i]) << offset
+		} else {
+			c += int(buf[i])
+		}
+	}
+
+	if c == 0 && len(buf) > FrameLengthFieldSize {
+		// skip the first byte and contine reading frame length from buf
+		buf = buf[1:]
+		return ReadFrameLength(buf)
+	}
+
+	return c
+}

pkg/framing/payload.go

@@ -0,0 +1,22 @@
+package framing
+
+// PayloadFrame represents a payload frame.
+type PayloadFrame struct {
+	data []byte
+}
+
+// NewPayloadFrame inits a new PayloadFrame.
+func NewPayloadFrame(data []byte) *PayloadFrame {
+	return &PayloadFrame{
+		data: data,
+	}
+}
+
+func (p *PayloadFrame) Bytes() []byte {
+	len := len(p.data)
+	buf := make([]byte, FrameLengthFieldSize)
+
+	appendFrameLength(buf, len)
+	buf = append(buf, p.data...)
+	return buf
+}

pkg/quic/conn.go

@@ -81,7 +81,6 @@ func (c *QuicConn) Healthcheck() {
 				if c.OnHeartbeatExpired != nil {
 					c.OnHeartbeatExpired()
 				}
-				break loop
 			}
 		}
 	}()

pkg/runtime/handler.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/reactivex/rxgo/v2"
 	"github.com/yomorun/yomo/pkg/client"
+	"github.com/yomorun/yomo/pkg/decoder"
 	"github.com/yomorun/yomo/pkg/quic"
 	"github.com/yomorun/yomo/pkg/serverless"
 )
@@ -158,13 +159,12 @@ func (s *quicHandler) receiveDataFromZipperSenders() {
 			}
 
 			go func() {
+				fd := decoder.NewFrameDecoder(receiver)
 				for {
-					buf := make([]byte, 3*1024)
-					n, err := receiver.Read(buf)
+					buf, err := fd.Read(false)
 					if err != nil {
 						break
 					} else {
-						value := buf[:n]
 						// send data to sinks
 						for _, sink := range sinks {
 							go func(sf serverless.GetSinkFunc, buf []byte) {
@@ -176,7 +176,7 @@ func (s *quicHandler) receiveDataFromZipperSenders() {
 										cancel()
 									}
 								}
-							}(sink, value)
+							}(sink, buf)
 						}
 					}
 				}

pkg/rx/rxstream_operator.go

@@ -57,14 +57,13 @@ func FromReader(reader io.Reader) RxStream {
 	go func() {
 		defer close(next)
 
+		fd := decoder.NewFrameDecoder(reader)
 		for {
-			buf := make([]byte, 3*1024)
-			n, err := reader.Read(buf)
+			buf, err := fd.Read(false)
 			if err != nil {
 				break
 			} else {
-				value := buf[:n]
-				next <- Of(value)
+				next <- Of(buf)
 			}
 		}
 	}()
@@ -105,16 +104,16 @@ func FromReaderWithFunc(f func() io.Reader) RxStream {
 			if reader == nil {
 				time.Sleep(time.Second)
 			} else {
-				buf := make([]byte, 3*1024)
-				n, err := reader.Read(buf)
-				if err != nil {
-					fmt.Println("[FromReaderWithFunc Reader Error] ", err)
-				} else {
-					value := buf[:n]
-					next <- Of(value)
+				fd := decoder.NewFrameDecoder(reader)
+				for {
+					buf, err := fd.Read(false)
+					if err != nil {
+						break
+					} else {
+						next <- Of(buf)
+					}
 				}
 			}
-
 		}
 	}()
 
@@ -658,13 +657,13 @@ func (s *RxStreamImpl) MergeReadWriterWithFunc(rwf serverless.GetFlowFunc, opts
 				if rw == nil {
 					time.Sleep(time.Second)
 				} else {
-					buf := make([]byte, 3*1024)
-					n, err := rw.Read(buf)
-
-					if err != nil {
+					fd := decoder.NewFrameDecoder(rw)
+					buf, err := fd.Read(false)
+					if err != nil && err != io.EOF {
 						cancel()
+						break
 					} else {
-						response <- buf[:n]
+						response <- buf
 					}
 				}
 			}