streaming rpc support (backend streams to the frontend) (#120)

pkg/wshrpc/wshclient/wshclient.go

@@ -13,62 +13,67 @@ import (
 
 // command "controller:input", wshserver.BlockInputCommand
 func BlockInputCommand(w *wshutil.WshRpc, data wshrpc.CommandBlockInputData, opts *wshrpc.WshRpcCommandOpts) error {
-    _, err := sendRpcRequestHelper[any](w, "controller:input", data, opts)
+    _, err := sendRpcRequestCallHelper[any](w, "controller:input", data, opts)
     return err
 }
 
 // command "controller:restart", wshserver.BlockRestartCommand
 func BlockRestartCommand(w *wshutil.WshRpc, data wshrpc.CommandBlockRestartData, opts *wshrpc.WshRpcCommandOpts) error {
-    _, err := sendRpcRequestHelper[any](w, "controller:restart", data, opts)
+    _, err := sendRpcRequestCallHelper[any](w, "controller:restart", data, opts)
     return err
 }
 
 // command "createblock", wshserver.CreateBlockCommand
 func CreateBlockCommand(w *wshutil.WshRpc, data wshrpc.CommandCreateBlockData, opts *wshrpc.WshRpcCommandOpts) (*waveobj.ORef, error) {
-    resp, err := sendRpcRequestHelper[*waveobj.ORef](w, "createblock", data, opts)
+    resp, err := sendRpcRequestCallHelper[*waveobj.ORef](w, "createblock", data, opts)
     return resp, err
 }
 
 // command "file:append", wshserver.AppendFileCommand
 func AppendFileCommand(w *wshutil.WshRpc, data wshrpc.CommandAppendFileData, opts *wshrpc.WshRpcCommandOpts) error {
-    _, err := sendRpcRequestHelper[any](w, "file:append", data, opts)
+    _, err := sendRpcRequestCallHelper[any](w, "file:append", data, opts)
     return err
 }
 
 // command "file:appendijson", wshserver.AppendIJsonCommand
 func AppendIJsonCommand(w *wshutil.WshRpc, data wshrpc.CommandAppendIJsonData, opts *wshrpc.WshRpcCommandOpts) error {
-    _, err := sendRpcRequestHelper[any](w, "file:appendijson", data, opts)
+    _, err := sendRpcRequestCallHelper[any](w, "file:appendijson", data, opts)
     return err
 }
 
 // command "getmeta", wshserver.GetMetaCommand
 func GetMetaCommand(w *wshutil.WshRpc, data wshrpc.CommandGetMetaData, opts *wshrpc.WshRpcCommandOpts) (map[string]interface {}, error) {
-    resp, err := sendRpcRequestHelper[map[string]interface {}](w, "getmeta", data, opts)
+    resp, err := sendRpcRequestCallHelper[map[string]interface {}](w, "getmeta", data, opts)
     return resp, err
 }
 
 // command "message", wshserver.MessageCommand
 func MessageCommand(w *wshutil.WshRpc, data wshrpc.CommandMessageData, opts *wshrpc.WshRpcCommandOpts) error {
-    _, err := sendRpcRequestHelper[any](w, "message", data, opts)
+    _, err := sendRpcRequestCallHelper[any](w, "message", data, opts)
     return err
 }
 
 // command "resolveids", wshserver.ResolveIdsCommand
 func ResolveIdsCommand(w *wshutil.WshRpc, data wshrpc.CommandResolveIdsData, opts *wshrpc.WshRpcCommandOpts) (wshrpc.CommandResolveIdsRtnData, error) {
-    resp, err := sendRpcRequestHelper[wshrpc.CommandResolveIdsRtnData](w, "resolveids", data, opts)
+    resp, err := sendRpcRequestCallHelper[wshrpc.CommandResolveIdsRtnData](w, "resolveids", data, opts)
     return resp, err
 }
 
 // command "setmeta", wshserver.SetMetaCommand
 func SetMetaCommand(w *wshutil.WshRpc, data wshrpc.CommandSetMetaData, opts *wshrpc.WshRpcCommandOpts) error {
-    _, err := sendRpcRequestHelper[any](w, "setmeta", data, opts)
+    _, err := sendRpcRequestCallHelper[any](w, "setmeta", data, opts)
     return err
 }
 
 // command "setview", wshserver.BlockSetViewCommand
 func BlockSetViewCommand(w *wshutil.WshRpc, data wshrpc.CommandBlockSetViewData, opts *wshrpc.WshRpcCommandOpts) error {
-    _, err := sendRpcRequestHelper[any](w, "setview", data, opts)
+    _, err := sendRpcRequestCallHelper[any](w, "setview", data, opts)
     return err
 }
 
+// command "streamtest", wshserver.RespStreamTest
+func RespStreamTest(w *wshutil.WshRpc, opts *wshrpc.WshRpcCommandOpts) chan wshrpc.RespOrErrorUnion[int] {
+    return sendRpcRequestResponseStreamHelper[int](w, "streamtest", nil, opts)
+}
+
 

pkg/wshrpc/wshclient/wshclientutil.go

@@ -9,7 +9,7 @@ import (
 	"github.com/wavetermdev/thenextwave/pkg/wshutil"
 )
 
-func sendRpcRequestHelper[T any](w *wshutil.WshRpc, command string, data interface{}, opts *wshrpc.WshRpcCommandOpts) (T, error) {
+func sendRpcRequestCallHelper[T any](w *wshutil.WshRpc, command string, data interface{}, opts *wshrpc.WshRpcCommandOpts) (T, error) {
 	var respData T
 	if opts.NoResponse {
 		err := w.SendCommand(command, data)
@@ -28,3 +28,36 @@ func sendRpcRequestHelper[T any](w *wshutil.WshRpc, command string, data interfa
 	}
 	return respData, nil
 }
+
+func sendRpcRequestResponseStreamHelper[T any](w *wshutil.WshRpc, command string, data interface{}, opts *wshrpc.WshRpcCommandOpts) chan wshrpc.RespOrErrorUnion[T] {
+	respChan := make(chan wshrpc.RespOrErrorUnion[T])
+	reqHandler, err := w.SendComplexRequest(command, data, true, opts.Timeout)
+	if err != nil {
+		go func() {
+			respChan <- wshrpc.RespOrErrorUnion[T]{Error: err}
+			close(respChan)
+		}()
+	} else {
+		go func() {
+			defer close(respChan)
+			for {
+				if reqHandler.ResponseDone() {
+					break
+				}
+				resp, err := reqHandler.NextResponse()
+				if err != nil {
+					respChan <- wshrpc.RespOrErrorUnion[T]{Error: err}
+					break
+				}
+				var respData T
+				err = utilfn.ReUnmarshal(&respData, resp)
+				if err != nil {
+					respChan <- wshrpc.RespOrErrorUnion[T]{Error: err}
+					break
+				}
+				respChan <- wshrpc.RespOrErrorUnion[T]{Response: respData}
+			}
+		}()
+	}
+	return respChan
+}

pkg/wshrpc/wshrpctypes.go

@@ -35,6 +35,11 @@ var DataTypeMap = map[string]reflect.Type{
 	"oref":          reflect.TypeOf(waveobj.ORef{}),
 }
 
+type RespOrErrorUnion[T any] struct {
+	Response T
+	Error    error
+}
+
 // for frontend
 type WshServerCommandMeta struct {
 	CommandType string `json:"commandtype"`

pkg/wshrpc/wshserver/wshserver.go

@@ -3,14 +3,14 @@
 
 package wshserver
 
+// this file contains the implementation of the wsh server methods
+
 import (
 	"context"
 	"encoding/base64"
 	"fmt"
 	"io/fs"
 	"log"
-	"net"
-	"os"
 	"reflect"
 	"strings"
 	"time"
@@ -18,33 +18,19 @@ import (
 	"github.com/wavetermdev/thenextwave/pkg/blockcontroller"
 	"github.com/wavetermdev/thenextwave/pkg/eventbus"
 	"github.com/wavetermdev/thenextwave/pkg/filestore"
-	"github.com/wavetermdev/thenextwave/pkg/util/utilfn"
-	"github.com/wavetermdev/thenextwave/pkg/wavebase"
 	"github.com/wavetermdev/thenextwave/pkg/waveobj"
 	"github.com/wavetermdev/thenextwave/pkg/wshrpc"
 	"github.com/wavetermdev/thenextwave/pkg/wshutil"
 	"github.com/wavetermdev/thenextwave/pkg/wstore"
 )
 
-const (
-	DefaultOutputChSize = 32
-	DefaultInputChSize  = 32
-)
-
-type WshServer struct{}
-
-var WshServerImpl = WshServer{}
-var contextRType = reflect.TypeOf((*context.Context)(nil)).Elem()
-
-type WshServerMethodDecl struct {
-	Command                 string
-	CommandType             string
-	MethodName              string
-	Method                  reflect.Value
-	CommandDataType         reflect.Type
-	DefaultResponseDataType reflect.Type
-	RequestDataTypes        []reflect.Type // for streaming requests
-	ResponseDataTypes       []reflect.Type // for streaming responses
+var RespStreamTest_MethodDecl = &WshServerMethodDecl{
+	Command:                 "streamtest",
+	CommandType:             wshutil.RpcType_ResponseStream,
+	MethodName:              "RespStreamTest",
+	Method:                  reflect.ValueOf(WshServerImpl.RespStreamTest),
+	CommandDataType:         nil,
+	DefaultResponseDataType: reflect.TypeOf((int)(0)),
 }
 
 var WshServerCommandToDeclMap = map[string]*WshServerMethodDecl{
@@ -58,37 +44,28 @@ var WshServerCommandToDeclMap = map[string]*WshServerMethodDecl{
 	wshrpc.Command_BlockInput:  GetWshServerMethod(wshrpc.Command_BlockInput, wshutil.RpcType_Call, "BlockInputCommand", WshServerImpl.BlockInputCommand),
 	wshrpc.Command_AppendFile:  GetWshServerMethod(wshrpc.Command_AppendFile, wshutil.RpcType_Call, "AppendFileCommand", WshServerImpl.AppendFileCommand),
 	wshrpc.Command_AppendIJson: GetWshServerMethod(wshrpc.Command_AppendIJson, wshutil.RpcType_Call, "AppendIJsonCommand", WshServerImpl.AppendIJsonCommand),
+	"streamtest":               RespStreamTest_MethodDecl,
 }
 
-func GetWshServerMethod(command string, commandType string, methodName string, methodFunc any) *WshServerMethodDecl {
-	methodVal := reflect.ValueOf(methodFunc)
-	methodType := methodVal.Type()
-	if methodType.Kind() != reflect.Func {
-		panic(fmt.Sprintf("methodVal must be a function got [%v]", methodType))
-	}
-	if methodType.In(0) != contextRType {
-		panic(fmt.Sprintf("methodVal must have a context as the first argument %v", methodType))
-	}
-	var defResponseType reflect.Type
-	if methodType.NumOut() > 1 {
-		defResponseType = methodType.Out(0)
-	}
-	rtn := &WshServerMethodDecl{
-		Command:                 command,
-		CommandType:             commandType,
-		MethodName:              methodName,
-		Method:                  methodVal,
-		CommandDataType:         methodType.In(1),
-		DefaultResponseDataType: defResponseType,
-	}
-	return rtn
-}
-
+// for testing
 func (ws *WshServer) MessageCommand(ctx context.Context, data wshrpc.CommandMessageData) error {
 	log.Printf("MESSAGE: %s | %q\n", data.ORef, data.Message)
 	return nil
 }
 
+// for testing
+func (ws *WshServer) RespStreamTest(ctx context.Context) chan wshrpc.RespOrErrorUnion[int] {
+	rtn := make(chan wshrpc.RespOrErrorUnion[int])
+	go func() {
+		for i := 1; i <= 5; i++ {
+			rtn <- wshrpc.RespOrErrorUnion[int]{Response: i}
+			time.Sleep(1 * time.Second)
+		}
+		close(rtn)
+	}()
+	return rtn
+}
+
 func (ws *WshServer) GetMetaCommand(ctx context.Context, data wshrpc.CommandGetMetaData) (wshrpc.MetaDataType, error) {
 	log.Printf("calling meta: %s\n", data.ORef)
 	obj, err := wstore.DBGetORef(ctx, data.ORef)
@@ -317,101 +294,3 @@ func (ws *WshServer) AppendIJsonCommand(ctx context.Context, data wshrpc.Command
 	})
 	return nil
 }
-
-func decodeRtnVals(rtnVals []reflect.Value) (any, error) {
-	switch len(rtnVals) {
-	case 0:
-		return nil, nil
-	case 1:
-		errIf := rtnVals[0].Interface()
-		if errIf == nil {
-			return nil, nil
-		}
-		return nil, errIf.(error)
-	case 2:
-		errIf := rtnVals[1].Interface()
-		if errIf == nil {
-			return rtnVals[0].Interface(), nil
-		}
-		return rtnVals[0].Interface(), errIf.(error)
-	default:
-		return nil, fmt.Errorf("too many return values: %d", len(rtnVals))
-	}
-}
-
-func mainWshServerHandler(handler *wshutil.RpcResponseHandler) {
-	command := handler.GetCommand()
-	methodDecl := WshServerCommandToDeclMap[command]
-	if methodDecl == nil {
-		handler.SendResponseError(fmt.Errorf("command %q not found", command))
-		return
-	}
-	var callParams []reflect.Value
-	callParams = append(callParams, reflect.ValueOf(handler.Context()))
-	if methodDecl.CommandDataType != nil {
-		commandData := reflect.New(methodDecl.CommandDataType).Interface()
-		err := utilfn.ReUnmarshal(commandData, handler.GetCommandRawData())
-		if err != nil {
-			handler.SendResponseError(fmt.Errorf("error re-marshalling command data: %w", err))
-			return
-		}
-		wshrpc.HackRpcContextIntoData(commandData, handler.GetRpcContext())
-		callParams = append(callParams, reflect.ValueOf(commandData).Elem())
-	}
-	rtnVals := methodDecl.Method.Call(callParams)
-	rtnData, rtnErr := decodeRtnVals(rtnVals)
-	if rtnErr != nil {
-		handler.SendResponseError(rtnErr)
-		return
-	} else {
-		handler.SendResponse(rtnData, true)
-	}
-}
-
-func MakeUnixListener(sockName string) (net.Listener, error) {
-	os.Remove(sockName) // ignore error
-	rtn, err := net.Listen("unix", sockName)
-	if err != nil {
-		return nil, fmt.Errorf("error creating listener at %v: %v", sockName, err)
-	}
-	os.Chmod(sockName, 0700)
-	log.Printf("Server listening on %s\n", sockName)
-	return rtn, nil
-}
-
-func runWshRpcWithStream(conn net.Conn) {
-	defer conn.Close()
-	inputCh := make(chan []byte, DefaultInputChSize)
-	outputCh := make(chan []byte, DefaultOutputChSize)
-	go wshutil.AdaptMsgChToStream(outputCh, conn)
-	go wshutil.AdaptStreamToMsgCh(conn, inputCh)
-	wshutil.MakeWshRpc(inputCh, outputCh, wshutil.RpcContext{}, mainWshServerHandler)
-}
-
-func RunWshRpcOverListener(listener net.Listener) {
-	go func() {
-		for {
-			conn, err := listener.Accept()
-			if err != nil {
-				log.Printf("error accepting connection: %v\n", err)
-				continue
-			}
-			go runWshRpcWithStream(conn)
-		}
-	}()
-}
-
-func RunDomainSocketWshServer() error {
-	sockName := wavebase.GetDomainSocketName()
-	listener, err := MakeUnixListener(sockName)
-	if err != nil {
-		return fmt.Errorf("error starging unix listener for wsh-server: %w", err)
-	}
-	defer listener.Close()
-	RunWshRpcOverListener(listener)
-	return nil
-}
-
-func MakeWshServer(inputCh chan []byte, outputCh chan []byte, initialCtx wshutil.RpcContext) {
-	wshutil.MakeWshRpc(inputCh, outputCh, initialCtx, mainWshServerHandler)
-}

pkg/wshrpc/wshserver/wshserverutil.go

@@ -0,0 +1,193 @@
+// Copyright 2024, Command Line Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package wshserver
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net"
+	"os"
+	"reflect"
+
+	"github.com/wavetermdev/thenextwave/pkg/util/utilfn"
+	"github.com/wavetermdev/thenextwave/pkg/wavebase"
+	"github.com/wavetermdev/thenextwave/pkg/wshrpc"
+	"github.com/wavetermdev/thenextwave/pkg/wshutil"
+)
+
+// this file contains the generic types and functions that create and power the WSH server
+
+const (
+	DefaultOutputChSize = 32
+	DefaultInputChSize  = 32
+)
+
+type WshServer struct{}
+
+type WshServerMethodDecl struct {
+	Command                 string
+	CommandType             string
+	MethodName              string
+	Method                  reflect.Value
+	CommandDataType         reflect.Type
+	DefaultResponseDataType reflect.Type
+	RequestDataTypes        []reflect.Type // for streaming requests
+	ResponseDataTypes       []reflect.Type // for streaming responses
+}
+
+var WshServerImpl = WshServer{}
+var contextRType = reflect.TypeOf((*context.Context)(nil)).Elem()
+
+func GetWshServerMethod(command string, commandType string, methodName string, methodFunc any) *WshServerMethodDecl {
+	methodVal := reflect.ValueOf(methodFunc)
+	methodType := methodVal.Type()
+	if methodType.Kind() != reflect.Func {
+		panic(fmt.Sprintf("methodVal must be a function got [%v]", methodType))
+	}
+	if methodType.In(0) != contextRType {
+		panic(fmt.Sprintf("methodVal must have a context as the first argument %v", methodType))
+	}
+	var defResponseType reflect.Type
+	if methodType.NumOut() > 1 {
+		defResponseType = methodType.Out(0)
+	}
+	rtn := &WshServerMethodDecl{
+		Command:                 command,
+		CommandType:             commandType,
+		MethodName:              methodName,
+		Method:                  methodVal,
+		CommandDataType:         methodType.In(1),
+		DefaultResponseDataType: defResponseType,
+	}
+	return rtn
+}
+
+func decodeRtnVals(rtnVals []reflect.Value) (any, error) {
+	switch len(rtnVals) {
+	case 0:
+		return nil, nil
+	case 1:
+		errIf := rtnVals[0].Interface()
+		if errIf == nil {
+			return nil, nil
+		}
+		return nil, errIf.(error)
+	case 2:
+		errIf := rtnVals[1].Interface()
+		if errIf == nil {
+			return rtnVals[0].Interface(), nil
+		}
+		return rtnVals[0].Interface(), errIf.(error)
+	default:
+		return nil, fmt.Errorf("too many return values: %d", len(rtnVals))
+	}
+}
+
+func mainWshServerHandler(handler *wshutil.RpcResponseHandler) bool {
+	command := handler.GetCommand()
+	methodDecl := WshServerCommandToDeclMap[command]
+	if methodDecl == nil {
+		handler.SendResponseError(fmt.Errorf("command %q not found", command))
+		return true
+	}
+	var callParams []reflect.Value
+	callParams = append(callParams, reflect.ValueOf(handler.Context()))
+	if methodDecl.CommandDataType != nil {
+		commandData := reflect.New(methodDecl.CommandDataType).Interface()
+		err := utilfn.ReUnmarshal(commandData, handler.GetCommandRawData())
+		if err != nil {
+			handler.SendResponseError(fmt.Errorf("error re-marshalling command data: %w", err))
+			return true
+		}
+		wshrpc.HackRpcContextIntoData(commandData, handler.GetRpcContext())
+		callParams = append(callParams, reflect.ValueOf(commandData).Elem())
+	}
+	if methodDecl.CommandType == wshutil.RpcType_Call {
+		rtnVals := methodDecl.Method.Call(callParams)
+		rtnData, rtnErr := decodeRtnVals(rtnVals)
+		if rtnErr != nil {
+			handler.SendResponseError(rtnErr)
+			return true
+		}
+		handler.SendResponse(rtnData, true)
+		return true
+	} else if methodDecl.CommandType == wshutil.RpcType_ResponseStream {
+		rtnVals := methodDecl.Method.Call(callParams)
+		rtnChVal := rtnVals[0]
+		if rtnChVal.IsNil() {
+			handler.SendResponse(nil, true)
+			return true
+		}
+		go func() {
+			defer handler.Finalize()
+			// must use reflection here because we don't know the generic type of RespOrErrorUnion
+			for {
+				respVal, ok := rtnChVal.Recv()
+				if !ok {
+					break
+				}
+				errorVal := respVal.FieldByName("Error")
+				if !errorVal.IsNil() {
+					handler.SendResponseError(errorVal.Interface().(error))
+					break
+				}
+				respData := respVal.FieldByName("Response").Interface()
+				handler.SendResponse(respData, false)
+			}
+		}()
+		return false
+	} else {
+		handler.SendResponseError(fmt.Errorf("unsupported command type %q", methodDecl.CommandType))
+		return true
+	}
+}
+
+func MakeUnixListener(sockName string) (net.Listener, error) {
+	os.Remove(sockName) // ignore error
+	rtn, err := net.Listen("unix", sockName)
+	if err != nil {
+		return nil, fmt.Errorf("error creating listener at %v: %v", sockName, err)
+	}
+	os.Chmod(sockName, 0700)
+	log.Printf("Server listening on %s\n", sockName)
+	return rtn, nil
+}
+
+func runWshRpcWithStream(conn net.Conn) {
+	defer conn.Close()
+	inputCh := make(chan []byte, DefaultInputChSize)
+	outputCh := make(chan []byte, DefaultOutputChSize)
+	go wshutil.AdaptMsgChToStream(outputCh, conn)
+	go wshutil.AdaptStreamToMsgCh(conn, inputCh)
+	wshutil.MakeWshRpc(inputCh, outputCh, wshutil.RpcContext{}, mainWshServerHandler)
+}
+
+func RunWshRpcOverListener(listener net.Listener) {
+	go func() {
+		for {
+			conn, err := listener.Accept()
+			if err != nil {
+				log.Printf("error accepting connection: %v\n", err)
+				continue
+			}
+			go runWshRpcWithStream(conn)
+		}
+	}()
+}
+
+func RunDomainSocketWshServer() error {
+	sockName := wavebase.GetDomainSocketName()
+	listener, err := MakeUnixListener(sockName)
+	if err != nil {
+		return fmt.Errorf("error starging unix listener for wsh-server: %w", err)
+	}
+	defer listener.Close()
+	RunWshRpcOverListener(listener)
+	return nil
+}
+
+func MakeWshServer(inputCh chan []byte, outputCh chan []byte, initialCtx wshutil.RpcContext) {
+	wshutil.MakeWshRpc(inputCh, outputCh, initialCtx, mainWshServerHandler)
+}