starpc 0.10.6 → 0.10.8

package/srpc/rwc-conn.go

@@ -0,0 +1,229 @@
+package srpc
+
+import (
+	"context"
+	"io"
+	"net"
+	"os"
+	"sync"
+	"time"
+)
+
+// connPktSize is the size of the buffers to use for packets for the RwcConn.
+const connPktSize = 2048
+
+// RwcConn implements a Conn with a buffered ReadWriteCloser.
+type RwcConn struct {
+	// ctx is the context
+	ctx context.Context
+	// ctxCancel is the context canceler
+	ctxCancel context.CancelFunc
+	// rwc is the read-write-closer
+	rwc io.ReadWriteCloser
+	// laddr is the local addr
+	laddr net.Addr
+	// raddr is the remote addr
+	raddr net.Addr
+
+	ar       sync.Pool   // packet arena
+	rd       time.Time   // read deadline
+	wd       time.Time   // write deadline
+	packetCh chan []byte // packet ch
+	closeErr error
+}
+
+// NewRwcConn constructs a new packet conn and starts the rx pump.
+func NewRwcConn(
+	ctx context.Context,
+	rwc io.ReadWriteCloser,
+	laddr, raddr net.Addr,
+	bufferPacketN int,
+) *RwcConn {
+	ctx, ctxCancel := context.WithCancel(ctx)
+	if bufferPacketN <= 0 {
+		bufferPacketN = 10
+	}
+
+	c := &RwcConn{
+		ctx:       ctx,
+		ctxCancel: ctxCancel,
+		rwc:       rwc,
+		laddr:     laddr,
+		raddr:     raddr,
+		packetCh:  make(chan []byte, bufferPacketN),
+	}
+	go func() {
+		_ = c.rxPump()
+	}()
+	return c
+}
+
+// LocalAddr returns the local network address.
+func (p *RwcConn) LocalAddr() net.Addr {
+	return p.laddr
+}
+
+// RemoteAddr returns the bound remote network address.
+func (p *RwcConn) RemoteAddr() net.Addr {
+	return p.raddr
+}
+
+// Read reads data from the connection.
+// Read can be made to time out and return an error after a fixed
+// time limit; see SetDeadline and SetReadDeadline.
+func (p *RwcConn) Read(b []byte) (n int, err error) {
+	deadline := p.rd
+	ctx := p.ctx
+	if !deadline.IsZero() {
+		var ctxCancel context.CancelFunc
+		ctx, ctxCancel = context.WithDeadline(ctx, deadline)
+		defer ctxCancel()
+	}
+
+	var pkt []byte
+	var ok bool
+	select {
+	case <-ctx.Done():
+		if !deadline.IsZero() {
+			return 0, os.ErrDeadlineExceeded
+		}
+		return 0, context.Canceled
+	case pkt, ok = <-p.packetCh:
+		if !ok {
+			err = p.closeErr
+			if err == nil {
+				err = io.EOF
+			}
+			return 0, err
+		}
+	}
+
+	pl := len(pkt)
+	copy(b, pkt)
+	p.ar.Put(&pkt)
+	if len(b) < pl {
+		return len(b), io.ErrShortBuffer
+	}
+	return pl, nil
+}
+
+// Write writes data to the connection.
+func (p *RwcConn) Write(pkt []byte) (n int, err error) {
+	if len(pkt) == 0 {
+		return 0, nil
+	}
+
+	written := 0
+	for written < len(pkt) {
+		n, err = p.rwc.Write(pkt[written:])
+		written += n
+		if err != nil {
+			return written, err
+		}
+	}
+	return written, nil
+}
+
+// SetDeadline sets the read and write deadlines associated
+// with the connection. It is equivalent to calling both
+// SetReadDeadline and SetWriteDeadline.
+//
+// A deadline is an absolute time after which I/O operations
+// fail instead of blocking. The deadline applies to all future
+// and pending I/O, not just the immediately following call to
+// Read or Write. After a deadline has been exceeded, the
+// connection can be refreshed by setting a deadline in the future.
+//
+// If the deadline is exceeded a call to Read or Write or to other
+// I/O methods will return an error that wraps os.ErrDeadlineExceeded.
+// This can be tested using errors.Is(err, os.ErrDeadlineExceeded).
+// The error's Timeout method will return true, but note that there
+// are other possible errors for which the Timeout method will
+// return true even if the deadline has not been exceeded.
+//
+// An idle timeout can be implemented by repeatedly extending
+// the deadline after successful ReadFrom or WriteTo calls.
+//
+// A zero value for t means I/O operations will not time out.
+func (p *RwcConn) SetDeadline(t time.Time) error {
+	p.rd = t
+	p.wd = t
+	return nil
+}
+
+// SetReadDeadline sets the deadline for future ReadFrom calls
+// and any currently-blocked ReadFrom call.
+// A zero value for t means ReadFrom will not time out.
+func (p *RwcConn) SetReadDeadline(t time.Time) error {
+	p.rd = t
+	return nil
+}
+
+// SetWriteDeadline sets the deadline for future WriteTo calls
+// and any currently-blocked WriteTo call.
+// Even if write times out, it may return n > 0, indicating that
+// some of the data was successfully written.
+// A zero value for t means WriteTo will not time out.
+func (p *RwcConn) SetWriteDeadline(t time.Time) error {
+	p.wd = t
+	return nil
+}
+
+// Close closes the connection.
+// Any blocked ReadFrom or WriteTo operations will be unblocked and return errors.
+func (p *RwcConn) Close() error {
+	return p.rwc.Close()
+}
+
+// getArenaBuf returns a buf from the packet arena with at least the given size
+func (p *RwcConn) getArenaBuf(size int) []byte {
+	var buf []byte
+	bufp := p.ar.Get()
+	if bufp != nil {
+		buf = *bufp.(*[]byte)
+	}
+	if size != 0 {
+		if cap(buf) < size {
+			buf = make([]byte, size)
+		} else {
+			buf = buf[:size]
+		}
+	} else {
+		buf = buf[:cap(buf)]
+	}
+	return buf
+}
+
+// rxPump receives messages from the underlying connection.
+func (p *RwcConn) rxPump() (rerr error) {
+	defer func() {
+		p.closeErr = rerr
+		close(p.packetCh)
+	}()
+
+	for {
+		select {
+		case <-p.ctx.Done():
+			return p.ctx.Err()
+		default:
+		}
+
+		pktBuf := p.getArenaBuf(int(connPktSize))
+		n, err := p.rwc.Read(pktBuf)
+		if n == 0 {
+			p.ar.Put(&pktBuf)
+		} else {
+			select {
+			case <-p.ctx.Done():
+				return context.Canceled
+			case p.packetCh <- pktBuf[:n]:
+			}
+		}
+		if err != nil {
+			return err
+		}
+	}
+}
+
+// _ is a type assertion
+var _ net.Conn = ((*RwcConn)(nil))

package/srpc/server-rpc.go

@@ -16,8 +16,8 @@ type ServerRPC struct {
 	ctxCancel context.CancelFunc
 	// writer is the writer to write messages to
 	writer Writer
-	// mux is the mux to handle calls
-	mux Mux
+	// invoker is the rpc call invoker
+	invoker Invoker
 	// service is the rpc service
 	service string
 	// method is the rpc method
@@ -36,10 +36,10 @@ type ServerRPC struct {
 
 // NewServerRPC constructs a new ServerRPC session.
 // note: call SetWriter before handling any incoming messages.
-func NewServerRPC(ctx context.Context, mux Mux) *ServerRPC {
+func NewServerRPC(ctx context.Context, invoker Invoker) *ServerRPC {
 	rpc := &ServerRPC{
-		dataCh: make(chan []byte, 5),
-		mux:    mux,
+		dataCh:  make(chan []byte, 5),
+		invoker: invoker,
 	}
 	rpc.ctx, rpc.ctxCancel = context.WithCancel(ctx)
 	return rpc
@@ -165,7 +165,7 @@ func (r *ServerRPC) invokeRPC() {
 	// ctx := r.ctx
 	serviceID, methodID := r.service, r.method
 	strm := NewMsgStream(r.ctx, r.writer, r.dataCh)
-	ok, err := r.mux.InvokeMethod(serviceID, methodID, strm)
+	ok, err := r.invoker.InvokeMethod(serviceID, methodID, strm)
 	if err == nil && !ok {
 		err = ErrUnimplemented
 	}

package/srpc/server.go

@@ -9,27 +9,27 @@ import (
 
 // Server handles incoming RPC streams with a mux.
 type Server struct {
-	// mux is the srpc mux
-	mux Mux
+	// invoker is the method invoker
+	invoker Invoker
 }
 
 // NewServer constructs a new SRPC server.
-func NewServer(mux Mux) *Server {
+func NewServer(invoker Invoker) *Server {
 	return &Server{
-		mux: mux,
+		invoker: invoker,
 	}
 }
 
-// GetMux returns the mux.
-func (s *Server) GetMux() Mux {
-	return s.mux
+// GetInvoker returns the invoker.
+func (s *Server) GetInvoker() Invoker {
+	return s.invoker
 }
 
 // HandleStream handles an incoming ReadWriteCloser stream.
 func (s *Server) HandleStream(ctx context.Context, rwc io.ReadWriteCloser) error {
 	subCtx, subCtxCancel := context.WithCancel(ctx)
 	defer subCtxCancel()
-	serverRPC := NewServerRPC(subCtx, s.mux)
+	serverRPC := NewServerRPC(subCtx, s.invoker)
 	prw := NewPacketReadWriter(rwc)
 	serverRPC.SetWriter(prw)
 	go prw.ReadPump(serverRPC.HandlePacket, serverRPC.HandleStreamClose)