p2p: add MsgPipe for protocol testing

This commit is contained in:
Felix Lange 2014-12-12 11:39:07 +01:00
parent e28c60caf9
commit 65e39bf20e
2 changed files with 140 additions and 0 deletions

View File

@ -3,9 +3,11 @@ package p2p
import ( import (
"bytes" "bytes"
"encoding/binary" "encoding/binary"
"errors"
"io" "io"
"io/ioutil" "io/ioutil"
"math/big" "math/big"
"sync/atomic"
"github.com/ethereum/go-ethereum/ethutil" "github.com/ethereum/go-ethereum/ethutil"
"github.com/ethereum/go-ethereum/rlp" "github.com/ethereum/go-ethereum/rlp"
@ -153,3 +155,78 @@ func (r *postrack) ReadByte() (byte, error) {
} }
return b, err return b, err
} }
// MsgPipe creates a message pipe. Reads on one end are matched
// with writes on the other. The pipe is full-duplex, both ends
// implement MsgReadWriter.
func MsgPipe() (*MsgPipeRW, *MsgPipeRW) {
var (
c1, c2 = make(chan Msg), make(chan Msg)
closing = make(chan struct{})
closed = new(int32)
rw1 = &MsgPipeRW{c1, c2, closing, closed}
rw2 = &MsgPipeRW{c2, c1, closing, closed}
)
return rw1, rw2
}
// ErrPipeClosed is returned from pipe operations after the
// pipe has been closed.
var ErrPipeClosed = errors.New("p2p: read or write on closed message pipe")
// MsgPipeRW is an endpoint of a MsgReadWriter pipe.
type MsgPipeRW struct {
w chan<- Msg
r <-chan Msg
closing chan struct{}
closed *int32
}
// WriteMsg sends a messsage on the pipe.
// It blocks until the receiver has consumed the message payload.
func (p *MsgPipeRW) WriteMsg(msg Msg) error {
if atomic.LoadInt32(p.closed) == 0 {
consumed := make(chan struct{}, 1)
msg.Payload = &eofSignal{msg.Payload, int64(msg.Size), consumed}
select {
case p.w <- msg:
if msg.Size > 0 {
// wait for payload read or discard
<-consumed
}
return nil
case <-p.closing:
}
}
return ErrPipeClosed
}
// EncodeMsg is a convenient shorthand for sending an RLP-encoded message.
func (p *MsgPipeRW) EncodeMsg(code uint64, data ...interface{}) error {
return p.WriteMsg(NewMsg(code, data...))
}
// ReadMsg returns a message sent on the other end of the pipe.
func (p *MsgPipeRW) ReadMsg() (Msg, error) {
if atomic.LoadInt32(p.closed) == 0 {
select {
case msg := <-p.r:
return msg, nil
case <-p.closing:
}
}
return Msg{}, ErrPipeClosed
}
// Close unblocks any pending ReadMsg and WriteMsg calls on both ends
// of the pipe. They will return ErrPipeClosed. Note that Close does
// not interrupt any reads from a message payload.
func (p *MsgPipeRW) Close() error {
if atomic.AddInt32(p.closed, 1) != 1 {
// someone else is already closing
atomic.StoreInt32(p.closed, 1) // avoid overflow
return nil
}
close(p.closing)
return nil
}

View File

@ -2,8 +2,11 @@ package p2p
import ( import (
"bytes" "bytes"
"fmt"
"io/ioutil" "io/ioutil"
"runtime"
"testing" "testing"
"time"
"github.com/ethereum/go-ethereum/ethutil" "github.com/ethereum/go-ethereum/ethutil"
) )
@ -68,3 +71,63 @@ func TestDecodeRealMsg(t *testing.T) {
t.Errorf("incorrect code %d, want %d", msg.Code, 0) t.Errorf("incorrect code %d, want %d", msg.Code, 0)
} }
} }
func ExampleMsgPipe() {
rw1, rw2 := MsgPipe()
go func() {
rw1.EncodeMsg(8, []byte{0, 0})
rw1.EncodeMsg(5, []byte{1, 1})
rw1.Close()
}()
for {
msg, err := rw2.ReadMsg()
if err != nil {
break
}
var data [1][]byte
msg.Decode(&data)
fmt.Printf("msg: %d, %x\n", msg.Code, data[0])
}
// Output:
// msg: 8, 0000
// msg: 5, 0101
}
func TestMsgPipeUnblockWrite(t *testing.T) {
loop:
for i := 0; i < 100; i++ {
rw1, rw2 := MsgPipe()
done := make(chan struct{})
go func() {
if err := rw1.EncodeMsg(1); err == nil {
t.Error("EncodeMsg returned nil error")
} else if err != ErrPipeClosed {
t.Error("EncodeMsg returned wrong error: got %v, want %v", err, ErrPipeClosed)
}
close(done)
}()
// this call should ensure that EncodeMsg is waiting to
// deliver sometimes. if this isn't done, Close is likely to
// be executed before EncodeMsg starts and then we won't test
// all the cases.
runtime.Gosched()
rw2.Close()
select {
case <-done:
case <-time.After(200 * time.Millisecond):
t.Errorf("write didn't unblock")
break loop
}
}
}
// This test should panic if concurrent close isn't implemented correctly.
func TestMsgPipeConcurrentClose(t *testing.T) {
rw1, _ := MsgPipe()
for i := 0; i < 10; i++ {
go rw1.Close()
}
}