refactor(runtime): migrate runtime codes from cli project to yomo (#189)

* refactor(runtime): migrate `runtime` related codes from cli project to yomo

* refactor(conn): add a common quic conn for client and server

* fix(mesh): retry to connect downstream zipper when the conn is disconnected

Author: xiaojian-hong

go.mod

@@ -9,6 +9,6 @@ require (
 	github.com/stretchr/objx v0.1.1 // indirect
 	github.com/yomorun/y3-codec-golang v1.6.9
 	golang.org/x/tools v0.0.0-20191125144606-a911d9008d1f // indirect
-	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v2 v2.4.0
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )

pkg/client/client.go

@@ -12,29 +12,6 @@ import (
 	"time"
 
 	"github.com/yomorun/yomo/pkg/quic"
-	"github.com/yomorun/yomo/pkg/rx"
-)
-
-const (
-	// ClientTypeSource represents the client type of Source.
-	ClientTypeSource = "source"
-
-	// ClientTypeServerless represents the client type of Serverless.
-	ClientTypeServerless = "serverless"
-
-	// ClientTypeZipperSender represents the client type of ZipperSender.
-	ClientTypeZipperSender = "zipper-sender"
-)
-
-var (
-	// SignalHeartbeat represents the signal of Heartbeat.
-	SignalHeartbeat = []byte{0}
-
-	// SignalAccepted represents the signal of Accpeted.
-	SignalAccepted = []byte{1}
-
-	// SignalFlowSink represents the signal for flow/sink.
-	SignalFlowSink = []byte{0, 0}
 )
 
 // NegotiationPayload represents the payload for negotiation.
@@ -56,69 +33,44 @@ type client interface {
 	RetryWithCount(count int) bool
 }
 
-// SourceClient is the client for YoMo-Source.
-// https://yomo.run/source
-type SourceClient interface {
-	client
-
-	// Connect to YoMo-Zipper
-	Connect(ip string, port int) (SourceClient, error)
-}
-
-// ServerlessClient is the client for YoMo-Serverless.
-type ServerlessClient interface {
-	client
-
-	// Connect to YoMo-Zipper
-	Connect(ip string, port int) (ServerlessClient, error)
-
-	// Pipe the Handler function.
-	Pipe(f func(rxstream rx.RxStream) rx.RxStream)
-}
-
-// ZipperSenderClient is the client for Zipper-Sender to connect the downsteam Zipper-Receiver  in edge-mesh.
-type ZipperSenderClient interface {
-	client
-
-	// Connect to downsteam Zipper-Receiver
-	Connect(ip string, port int) (ZipperSenderClient, error)
-}
-
 type clientImpl struct {
+	conn       *quic.QuicConn
 	zipperIP   string
 	zipperPort int
-	name       string
-	clientType string
 	readers    chan io.Reader
 	writer     io.Writer
 	session    quic.Client
-	signal     quic.Stream
-	stream     quic.Stream
-	heartbeat  chan byte
 	once       *sync.Once
 }
 
-type sourceClientImpl struct {
-	*clientImpl
-}
-
-type serverlessClientImpl struct {
-	*clientImpl
-}
-
-type zipperSenderClientImpl struct {
-	*clientImpl
-}
-
 // newClient creates a new client.
 func newClient(appName string, clientType string) *clientImpl {
 	c := &clientImpl{
-		name:       appName,
-		clientType: clientType,
-		readers:    make(chan io.Reader, 1),
-		heartbeat:  make(chan byte),
-		once:       new(sync.Once),
+		conn:    quic.NewQuicConn(appName, clientType),
+		readers: make(chan io.Reader, 1),
+		once:    new(sync.Once),
+	}
+
+	c.conn.OnHeartbeatReceived = func() {
+		// when the client received the heartbeat from server, send it back back to server (ping/pong).
+		c.conn.SendSignal(quic.SignalHeartbeat)
 	}
+
+	c.conn.OnHeartbeatExpired = func() {
+		c.once.Do(func() {
+			// reset stream to nil.
+			c.conn.Stream = nil
+
+			// reconnect when the heartbeat is expired.
+			c.connect(c.zipperIP, c.zipperPort)
+
+			// reset the sync.Once after 5s.
+			time.AfterFunc(5*time.Second, func() {
+				c.once = new(sync.Once)
+			})
+		})
+	}
+
 	return c
 }
 
@@ -129,65 +81,78 @@ func (c *clientImpl) connect(ip string, port int) (*clientImpl, error) {
 	c.zipperPort = port
 	addr := fmt.Sprintf("%s:%d", ip, port)
 	log.Println("Connecting to zipper", addr, "...")
+
 	// connect to yomo-zipper
 	quic_cli, err := quic.NewClient(addr)
 	if err != nil {
 		fmt.Println("client [NewClient] Error:", err)
 		return c, err
 	}
+
 	// create stream
 	quic_stream, err := quic_cli.CreateStream(context.Background())
 	if err != nil {
 		fmt.Println("client [CreateStream] Error:", err)
 		return c, err
 	}
 
+	// set session and signal
 	c.session = quic_cli
-	c.signal = quic_stream
+	c.conn.Signal = quic_stream
 
 	// send negotiation payload to zipper
 	payload := NegotiationPayload{
-		AppName:    c.name,
-		ClientType: c.clientType,
+		AppName:    c.conn.Name,
+		ClientType: c.conn.Type,
 	}
 	buf, _ := json.Marshal(payload)
-	_, err = c.signal.Write(buf)
+	err = c.conn.SendSignal(buf)
 
 	if err != nil {
 		fmt.Println("client [Write] Error:", err)
 		return c, err
 	}
 
-	// flow, sink create stream or heartbeat
 	accepted := make(chan bool)
 
+	c.handleSignal(accepted)
+	c.conn.Healthcheck()
+
+	// waiting when the connection is accepted.
+	<-accepted
+	log.Print("✅ Connected to zipper ", addr)
+	return c, nil
+}
+
+// handleSignal handles the logic when receiving signal from server.
+func (c *clientImpl) handleSignal(accepted chan bool) {
 	go func() {
 		defer close(accepted)
 		for {
 			buf := make([]byte, 2)
-			n, err := c.signal.Read(buf)
+			n, err := c.conn.Signal.Read(buf)
 			if err != nil {
 				break
 			}
 			value := buf[:n]
 
-			if bytes.Equal(value, SignalHeartbeat) {
+			if bytes.Equal(value, quic.SignalHeartbeat) {
 				// heartbeart
-				c.heartbeat <- buf[0]
-			} else if bytes.Equal(value, SignalAccepted) {
+				c.conn.Heartbeat <- buf[0]
+			} else if bytes.Equal(value, quic.SignalAccepted) {
 				// accepted
-				if c.clientType == ClientTypeSource || c.clientType == ClientTypeZipperSender {
-					// create the stream from source.
+				if c.conn.Type == quic.ConnTypeSource || c.conn.Type == quic.ConnTypeZipperSender {
+					// create stream for source.
 					stream, err := c.session.CreateStream(context.Background())
 					if err != nil {
 						fmt.Println("client [session.CreateStream] Error:", err)
 						break
 					}
-					c.stream = stream
+					c.conn.Stream = stream
 				}
 				accepted <- true
-			} else if bytes.Equal(value, SignalFlowSink) {
-				// create stream
+			} else if bytes.Equal(value, quic.SignalFlowSink) {
+				// create stream for flow/sink.
 				stream, err := c.session.CreateStream(context.Background())
 
 				if err != nil {
@@ -197,49 +162,16 @@ func (c *clientImpl) connect(ip string, port int) (*clientImpl, error) {
 
 				c.readers <- stream
 				c.writer = stream
-				stream.Write(SignalHeartbeat)
-			}
-		}
-	}()
-
-	go func() {
-		defer c.Close()
-		for {
-			select {
-			case item, ok := <-c.heartbeat:
-				if !ok {
-					return
-				}
-				_, err := c.signal.Write([]byte{item})
-
-				if err != nil {
-					return
-				}
-			case <-time.After(5 * time.Second):
-				// reconnect if didn't receive the heartbeat after 5s.
-				c.once.Do(func() {
-					c.connect(c.zipperIP, c.zipperPort)
-
-					// reset the sync.Once after 5s.
-					time.AfterFunc(5*time.Second, func() {
-						c.once = new(sync.Once)
-					})
-				})
+				stream.Write(quic.SignalHeartbeat)
 			}
 		}
-
 	}()
-
-	// waiting when the connection is accepted.
-	<-accepted
-	log.Print("✅ Connected to zipper ", addr)
-	return c, nil
 }
 
 // Write the data to downstream.
 func (c *clientImpl) Write(b []byte) (int, error) {
-	if c.stream != nil {
-		return c.stream.Write(b)
+	if c.conn.Stream != nil {
+		return c.conn.Stream.Write(b)
 	} else {
 		return 0, errors.New("not found stream")
 	}
@@ -273,85 +205,7 @@ func (c *clientImpl) RetryWithCount(count int) bool {
 // Close the client.
 func (c *clientImpl) Close() error {
 	err := c.session.Close()
-	c.heartbeat = make(chan byte)
-	c.signal = nil
+	c.conn.Heartbeat = make(chan byte)
+	c.conn.Signal = nil
 	return err
 }
-
-// NewSource setups the client of YoMo-Source.
-func NewSource(appName string) SourceClient {
-	c := &sourceClientImpl{
-		clientImpl: newClient(appName, ClientTypeSource),
-	}
-	return c
-}
-
-// Connect to yomo-zipper.
-func (c *sourceClientImpl) Connect(ip string, port int) (SourceClient, error) {
-	cli, err := c.connect(ip, port)
-	return &sourceClientImpl{
-		cli,
-	}, err
-}
-
-// NewServerless setups the client of YoMo-Serverless.
-// The "appName" should match the name of flows (or sinks) in workflow.yaml in zipper.
-func NewServerless(appName string) ServerlessClient {
-	c := &serverlessClientImpl{
-		clientImpl: newClient(appName, ClientTypeServerless),
-	}
-	return c
-}
-
-// Connect to yomo-zipper.
-func (c *serverlessClientImpl) Connect(ip string, port int) (ServerlessClient, error) {
-	cli, err := c.connect(ip, port)
-	return &serverlessClientImpl{
-		cli,
-	}, err
-}
-
-// Pipe the handler function in flow/sink serverless.
-func (c *serverlessClientImpl) Pipe(f func(rxstream rx.RxStream) rx.RxStream) {
-	rxstream := rx.FromReaderWithDecoder(c.readers)
-	stream := f(rxstream)
-
-	rxstream.Connect(context.Background())
-
-	for customer := range stream.Observe() {
-		if customer.Error() {
-			panic(customer.E)
-		} else if customer.V != nil {
-			if c.writer == nil {
-				continue
-			}
-
-			buf, ok := (customer.V).([]byte)
-			if !ok {
-				log.Print("❌ Please add the encode/marshal operator in the end of your Serverless handler.")
-				continue
-			}
-			_, err := c.writer.Write(buf)
-			if err != nil {
-				log.Print("❌ Send data to zipper failed. ", err)
-			}
-		}
-
-	}
-}
-
-// NewZipperSender setups the client of Zipper-Sender.
-func NewZipperSender(appName string) ZipperSenderClient {
-	c := &zipperSenderClientImpl{
-		clientImpl: newClient(appName, ClientTypeZipperSender),
-	}
-	return c
-}
-
-// Connect to downstream zipper-receiver in edge-mesh.
-func (c *zipperSenderClientImpl) Connect(ip string, port int) (ZipperSenderClient, error) {
-	cli, err := c.connect(ip, port)
-	return &zipperSenderClientImpl{
-		cli,
-	}, err
-}