tests passed

pkg/jsonrpc/message.go

+package jsonrpc
+
+import (
+	"encoding/json"
+	"io"
+
+	"github.com/go-faster/jx"
+)
+
+// MessageStream is a writer used to write jsonrpc message to a stream
+type MessageStream struct {
+	w  io.Writer
+	jx *jx.Writer
+}
+
+func NewStream(w io.Writer) (*MessageStream, error) {
+	enc := jx.GetWriter()
+	defer jx.PutWriter(enc)
+	enc.Grow(4096)
+	enc.ResetWriter(w)
+	enc.ObjStart()
+	enc.FieldStart("jsonrpc")
+	enc.Str("2.0")
+	enc.Close()
+	return &MessageStream{
+		w:  w,
+		jx: enc,
+	}, nil
+}
+
+func (m *MessageStream) Field(name string, value json.RawMessage) error {
+	m.jx.ResetWriter(m.w)
+	m.jx.Comma()
+	m.jx.FieldStart(name)
+	m.jx.Raw(value)
+	return m.jx.Close()
+}
+
+// Result returns a writecloser that writes to a result field
+func (m *MessageStream) Result() (io.Writer, error) {
+	m.jx.ResetWriter(m.w)
+	m.jx.Comma()
+	m.jx.FieldStart("result")
+	m.jx.Close()
+	return &MessageWriter{w: m.w}, nil
+}
+
+func (m *MessageStream) Close() error {
+	_, err := m.w.Write([]byte("}"))
+	return err
+}
+
+type MessageWriter struct {
+	w io.Writer
+}
+
+func (m *MessageWriter) Write(p []byte) (n int, err error) {
+	return m.w.Write(p)
+}

pkg/server/server.go

@@ -294,63 +294,50 @@ type callEnv struct {
 
 func (c *callResponder) send(ctx context.Context, env *callEnv) (err error) {
 	w := c.remote
-	enc := jx.GetWriter()
-	defer jx.PutWriter(enc)
-	enc.Grow(4096)
-	enc.ResetWriter(w)
-	enc.ObjStart()
-	enc.FieldStart("jsonrpc")
-	enc.Str("2.0")
+	s, err := jsonrpc.NewStream(w)
+	if err != nil {
+		return err
+	}
+	defer s.Close()
 	if env.id != nil {
-		enc.Comma()
-		enc.FieldStart("id")
-		enc.Raw(env.id.RawMessage())
+		s.Field("id", env.id.RawMessage())
 	}
 	if env.err != nil {
-		enc.Comma()
-		enc.FieldStart("error")
-		enc.Raw(jsonrpc.MarshalError(env.err))
-	} else {
+		s.Field("error", jsonrpc.MarshalError(env.err))
+		return nil
+	}
 	// if there is no error, we try to marshal the result
-		enc.Comma()
-		enc.FieldStart("result")
-		enc.Close()
-		enc.ResetWriter(w)
-		if env.v != nil {
+	wr, err := s.Result()
+	if err != nil {
+		return err
+	}
+	if env.v == nil {
+		_, err := wr.Write(jsonrpc.Null)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
 	switch cast := (env.v).(type) {
 	case json.RawMessage:
 		if len(cast) == 0 {
-					enc.Null()
 		} else {
-					enc.Raw(cast)
+			_, err := wr.Write(cast)
+			if err != nil {
+				return err
+			}
 		}
 	case *io.PipeReader:
-				_, err := io.Copy(w, cast)
+		_, err := io.Copy(wr, cast)
 		if err != nil {
 			return err
 		}
 		cast.Close()
 	case func(e io.Writer) error:
-				err = cast(w)
-			case func(e *jx.Writer) error:
-				err = cast(enc)
+		err = cast(wr)
 	default:
 		err = jjson.Encode(w, cast)
 	}
-		} else {
-			enc.Null()
-		}
-	}
-	if env.err == nil && err != nil {
-		enc.Comma()
-		enc.FieldStart("error")
-		enc.Raw(jsonrpc.MarshalError(err))
-	}
-	enc.ObjEnd()
-	err = enc.Close()
-	if err != nil {
-		return err
-	}
 	return nil
 }