mediocre-go-lib/mrpc/jstreamrpc/jstreamrpc.go

// Package jstreamrpc implements the mrpc interfaces using the jstream protocol.
// This implementation makes a few design decisions which are not specified in
// the protocol, but which are good best practices none-the-less.
package jstreamrpc

import (
	"context"
	"errors"
	"net"

	"github.com/mediocregopher/mediocre-go-lib/jstream"
	"github.com/mediocregopher/mediocre-go-lib/mrpc"
)

// TODO Error?
// TODO it'd be nice if the types here played nice with mrpc.ReflectClient

type debug struct {
	Debug mrpc.Debug `json:"debug,omitempty"`
}

type reqHead struct {
	debug
	Method string `json:"method"`
}

type resTail struct {
	debug
	Error error `json:"err,omitempty"`
}

type ctxVal int

const (
	ctxValR ctxVal = iota
	ctxValW
)

func unmarshalBody(i interface{}, el jstream.Element) error {
	switch iT := i.(type) {
	case *jstream.Element:
		*iT = el
		return nil
	default:
		return el.DecodeValue(i)
	}
}

func marshalBody(w *jstream.StreamWriter, i interface{}) error {
	switch iT := i.(type) {
	case func(*jstream.StreamWriter) error:
		return iT(w)
	default:
		return w.EncodeValue(iT)
	}
}

// HandleCall TODO
//
// If this returns an error then both r and w should be discarded and no longer
// used.
func HandleCall(
	ctx context.Context,
	r *jstream.StreamReader,
	w *jstream.StreamWriter,
	h mrpc.Handler,
) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	var head reqHead
	if err := r.Next().DecodeValue(&head); err != nil {
		return err
	} else if head.Method == "" {
		return errors.New("request head missing 'method' field")
	}

	ctx = context.WithValue(ctx, ctxValR, r)
	ctx = context.WithValue(ctx, ctxValW, w)

	body := r.Next()
	if body.Err != nil {
		return body.Err
	}

	rw := new(mrpc.ResponseWriter)
	h.ServeRPC(mrpc.Request{
		Context: ctx,
		Method:  head.Method,
		Unmarshal: func(i interface{}) error {
			return unmarshalBody(i, body)
		},
		Debug: head.debug.Debug,
	}, rw)

	resErr, resErrOk := rw.Response.(error)
	if resErrOk {
		if err := w.EncodeValue(nil); err != nil {
			return err
		}
	} else {
		if err := marshalBody(w, rw.Response); err != nil {
			if _, ok := err.(net.Error); ok {
				return err
			}
			resErr = err
		}
	}

	if err := w.EncodeValue(resTail{
		debug: debug{Debug: rw.Debug},
		Error: resErr,
	}); err != nil {
		return err
	}

	// make sure the body has been consumed before returning
	if err := body.Discard(); err != nil {
		return err
	}

	return nil
}

/*
func sqr(r mrpc.Request, rw *mrpc.ResponseWriter) {
	var inCh chan jstream.Element
	if err := r.Unmarshal(&inCh); err != nil {
		rw.Response = err
		return
	}

	outCh := make(chan int)
	go func() {
		defer close(outCh)
		for el := range inCh {
			var i int
			// TODO this is the problem right here. Getting this error out of
			// here so that the response stream can be canceled
			if err := el.DecodeValue(&i); err != nil {
				panic(err)
			}
			outCh <- i*i
		}
	}()

	rw.Response = outCh
}
*/