wip: mplex and mplex tests

This commit is contained in:
Dmitriy Ryajov 2019-09-03 19:42:00 -06:00
parent cc26d4e040
commit a0aaa29739
2 changed files with 98 additions and 50 deletions

View File

@ -7,11 +7,11 @@
## This file may not be copied, modified, or distributed except according to ## This file may not be copied, modified, or distributed except according to
## those terms. ## those terms.
import sequtils import tables, sequtils
import chronos import chronos
import ../varint, ../connection, import ../varint, ../connection,
../vbuffer, ../protocol, ../vbuffer, ../protocol,
../stream/bufferstream, ../stream/bufferstream, ../stream/lpstream,
muxer muxer
const MaxMsgSize* = 1 shl 20 # 1mb const MaxMsgSize* = 1 shl 20 # 1mb
@ -31,8 +31,8 @@ type
StreamHandler = proc(conn: Connection): Future[void] {.gcsafe.} StreamHandler = proc(conn: Connection): Future[void] {.gcsafe.}
Mplex* = ref object of Muxer Mplex* = ref object of Muxer
remote*: seq[Channel] remote*: Table[int, Channel]
local*: seq[Channel] local*: Table[int, Channel]
currentId*: int currentId*: int
maxChannels*: uint maxChannels*: uint
streamHandler*: StreamHandler streamHandler*: StreamHandler
@ -67,7 +67,7 @@ proc readHeader*(conn: Connection): Future[(uint, MessageType)] {.async, gcsafe.
if res != VarintStatus.Success: if res != VarintStatus.Success:
buffer.setLen(0) buffer.setLen(0)
return return
except TransportIncompleteError: except LPStreamIncompleteError:
buffer.setLen(0) buffer.setLen(0)
proc writeHeader*(conn: Connection, proc writeHeader*(conn: Connection,
@ -77,7 +77,8 @@ proc writeHeader*(conn: Connection,
## write lenght prefixed ## write lenght prefixed
var buf = initVBuffer() var buf = initVBuffer()
buf.writeVarint(LPSomeUVarint(id.uint shl 3 or msgType.uint)) buf.writeVarint(LPSomeUVarint(id.uint shl 3 or msgType.uint))
buf.writeVarint(LPSomeUVarint(size.uint)) if size > 0:
buf.writeVarint(LPSomeUVarint(size.uint))
buf.finish() buf.finish()
result = conn.write(buf.buffer) result = conn.write(buf.buffer)
@ -94,55 +95,64 @@ proc newChannel*(mplex: Mplex,
result.id = id result.id = id
result.mplex = mplex result.mplex = mplex
result.initiator = initiator result.initiator = initiator
result.writeHandler = handler result.initBufferStream(handler, size)
result.maxSize = size
proc closed*(s: Channel): bool = s.closedLocal and s.closedRemote proc closed*(s: Channel): bool = s.closedLocal and s.closedRemote
proc close*(s: Channel) {.async.} = discard proc close*(s: Channel) {.async.} = discard
proc reset*(s: Channel) {.async.} = discard proc reset*(s: Channel) {.async.} = discard
########################################## ##########################################
##
## Mplex ## Mplex
##
########################################## ##########################################
proc getChannelList(m: Mplex, initiator: bool): var seq[Channel] = proc getChannelList(m: Mplex, initiator: bool): var Table[int, Channel] =
if initiator: if initiator:
result = m.remote result = m.remote
else: else:
result = m.local result = m.local
proc newStream*(m: Mplex, proc newStreamInternal*(m: Mplex,
chanId: int = -1, initiator: bool = true,
initiator: bool = true): chanId: int):
Future[Connection] {.async, gcsafe.} = Future[Channel] {.async, gcsafe.} =
## create new channel/stream ## create new channel/stream
defer: inc(m.currentId) let id = if initiator: m.currentId.inc(); m.currentId else: chanId
let id = if chanId > -1: chanId else: m.currentId
proc writeHandler(data: seq[byte]): Future[void] {.async, gcsafe.} = proc writeHandler(data: seq[byte]): Future[void] {.async, gcsafe.} =
let msgType = if initiator: MessageType.MsgIn else: MessageType.MsgOut let msgType = if initiator: MessageType.MsgOut else: MessageType.MsgIn
await m.connection.writeHeader(id, msgType, data.len) # write header await m.connection.writeHeader(id, msgType, data.len) # write header
await m.connection.write(data) # write data await m.connection.write(data) # write data
let channel = newChannel(m, id, initiator, writeHandler) result = newChannel(m, id, initiator, writeHandler)
m.getChannelList(initiator)[id] = channel m.getChannelList(initiator)[id] = result
result = newConnection(channel)
proc handle*(m: Mplex) {.async, gcsafe.} = proc newStreamInternal*(m: Mplex): Future[Channel] {.gcsafe.} =
while not m.connection.closed: result = m.newStreamInternal(true, 0)
let (id, msgType) = await m.connection.readHeader()
let initiator = bool(ord(msgType) and 1) proc handle*(m: Mplex): Future[void] {.async, gcsafe.} =
case msgType: try:
of MessageType.New: while not m.connection.closed:
await m.streamHandler(await m.newStream(id.int, false)) let (id, msgType) = await m.connection.readHeader()
of MessageType.MsgIn, MessageType.MsgOut: let initiator = bool(ord(msgType) and 1)
await m.getChannelList(initiator)[id.int].pushTo(await m.connection.readLp()) case msgType:
of MessageType.CloseIn, MessageType.CloseOut: of MessageType.New:
await m.getChannelList(initiator)[id.int].close() let channel = await m.newStreamInternal(false, id.int)
of MessageType.ResetIn, MessageType.ResetOut: await m.streamHandler(newConnection(channel))
await m.getChannelList(initiator)[id.int].reset() of MessageType.MsgIn, MessageType.MsgOut:
else: raise newMplexUnknownMsgError() let channel = m.getChannelList(initiator)[id.int]
let msg = await m.connection.readLp()
await channel.pushTo(msg)
of MessageType.CloseIn, MessageType.CloseOut:
let channel = m.getChannelList(initiator)[id.int]
await channel.close()
of MessageType.ResetIn, MessageType.ResetOut:
let channel = m.getChannelList(initiator)[id.int]
await channel.reset()
else: raise newMplexUnknownMsgError()
except Exception as exc:
#TODO: add proper loging
discard
finally:
await m.connection.close()
proc newMplex*(conn: Connection, proc newMplex*(conn: Connection,
streamHandler: StreamHandler, streamHandler: StreamHandler,
@ -151,12 +161,15 @@ proc newMplex*(conn: Connection,
result.connection = conn result.connection = conn
result.maxChannels = maxChanns result.maxChannels = maxChanns
result.streamHandler = streamHandler result.streamHandler = streamHandler
result.remote = initTable[int, Channel]()
result.local = initTable[int, Channel]()
method newStream*(m: Mplex): Future[Connection] {.gcsafe.} = method newStream*(m: Mplex): Future[Connection] {.async, gcsafe.} =
result = m.newStream(true) let channel = await m.newStreamInternal()
await m.connection.writeHeader(channel.id, MessageType.New, 0)
result = newConnection(channel)
method close(m: Mplex) {.async, gcsafe.} = method close*(m: Mplex) {.async, gcsafe.} =
let futs = @[allFutures(m.remote.mapIt(it.close())), await allFutures(@[allFutures(toSeq(m.remote.values).mapIt(it.close())),
allFutures(m.local.mapIt(it.close()))] allFutures(toSeq(m.local.values).mapIt(it.close()))])
await allFutures(futs)
await m.connection.close() await m.connection.close()

View File

@ -43,7 +43,7 @@ suite "Mplex":
# check msg == fromHex("880102") # check msg == fromHex("880102")
# let conn = newConnection(newTestEncodeStream(encHandler)) # let conn = newConnection(newTestEncodeStream(encHandler))
# await conn.writeHeader(uint(17), MessageType.New, 2) # await conn.writeHeader(17, MessageType.New, 2)
# result = true # result = true
# check: # check:
@ -63,15 +63,16 @@ suite "Mplex":
# check: # check:
# waitFor(testDecodeHeader()) == true # waitFor(testDecodeHeader()) == true
test "e2e - new stream": test "e2e - read/write initiator":
proc testNewStream(): Future[bool] {.async.} = proc testNewStream(): Future[bool] {.async.} =
let ma: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53351") let ma: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53380")
proc connHandler(conn: Connection) {.async, gcsafe.} = proc connHandler(conn: Connection) {.async, gcsafe.} =
proc handleListen(stream: Connection) {.async, gcsafe.} = proc handleMplexListen(stream: Connection) {.async, gcsafe.} =
await stream.writeLp("Hello from stream!") await stream.writeLp("Hello from stream!")
await stream.close()
let mplexListen = newMplex(conn, handleListen) let mplexListen = newMplex(conn, handleMplexListen)
await mplexListen.handle() await mplexListen.handle()
let transport1: TcpTransport = newTransport(TcpTransport) let transport1: TcpTransport = newTransport(TcpTransport)
@ -79,11 +80,45 @@ suite "Mplex":
let transport2: TcpTransport = newTransport(TcpTransport) let transport2: TcpTransport = newTransport(TcpTransport)
let conn = await transport2.dial(ma) let conn = await transport2.dial(ma)
proc handleDial(stream: Connection) {.async, gcsafe.} =
let msg = await stream.readLp()
proc handleDial(stream: Connection) {.async, gcsafe.} = discard
let mplexDial = newMplex(conn, handleDial) let mplexDial = newMplex(conn, handleDial)
let handleFut = mplexDial.handle() let dialFut = mplexDial.handle()
let stream = await mplexDial.newStream()
check cast[string](await stream.readLp()) == "Hello from stream!"
await conn.close()
await dialFut
result = true
check:
waitFor(testNewStream()) == true
test "e2e - read/write receiver":
proc testNewStream(): Future[bool] {.async.} =
let ma: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53381")
proc connHandler(conn: Connection) {.async, gcsafe.} =
proc handleMplexListen(stream: Connection) {.async, gcsafe.} =
let msg = await stream.readLp()
check cast[string](msg) == "Hello from stream!"
await stream.close()
let mplexListen = newMplex(conn, handleMplexListen)
await mplexListen.handle()
let transport1: TcpTransport = newTransport(TcpTransport)
await transport1.listen(ma, connHandler)
let transport2: TcpTransport = newTransport(TcpTransport)
let conn = await transport2.dial(ma)
proc handleDial(stream: Connection) {.async, gcsafe.} = discard
let mplexDial = newMplex(conn, handleDial)
let dialFut = mplexDial.handle()
let stream = await mplexDial.newStream()
await stream.writeLp("Hello from stream!")
await dialFut
result = true result = true
check: check: