Added templates write(string), write(seq[byte])

Added consume() Added tests for write(string), write(seq[byte]), consume().
2018-06-04 19:42:54 +03:00 · 2018-06-04 19:42:54 +03:00 · 2b8eeef7aa
parent f4d4d57ccf
commit 2b8eeef7aa
2 changed files with 150 additions and 13 deletions
--- a/asyncdispatch2/transports/stream.nim
+++ b/asyncdispatch2/transports/stream.nim
@ -115,7 +115,6 @@ template checkPending(t: untyped) =
    raise newException(TransportError, "Read operation already pending!")

 template shiftBuffer(t, c: untyped) =
-  # ZAH: Nim is not C, you don't need to put () around template parameters
  if (t).offset > c:
    moveMem(addr((t).buffer[0]), addr((t).buffer[(c)]), (t).offset - (c))
    (t).offset = (t).offset - (c)
@ -812,6 +811,14 @@ proc write*(transp: StreamTransport, pbytes: pointer,
    raise transp.getError()
  result = nbytes

+template write*(transp: StreamTransport, msg: var string): untyped =
+  ## Write string ``msg`` using transport ``transp``.
+  write(transp, addr msg[0], len(msg))
+
+template write*(transp: StreamTransport, msg: var seq[byte]): untyped =
+  ## Write seq[byte] ``msg`` using transport ``transp``.
+  write(transp, addr msg[0], len(msg))
+
 proc writeFile*(transp: StreamTransport, handle: int,
                offset: uint = 0,
                size: int = 0): Future[void] {.async.} =
@ -833,6 +840,11 @@ proc writeFile*(transp: StreamTransport, handle: int,
  if WriteError in transp.state:
    raise transp.getError()

+proc atEof*(transp: StreamTransport): bool {.inline.} =
+  ## Returns ``true`` if ``transp`` is at EOF.
+  result = (transp.offset == 0) and (ReadEof in transp.state) and
+           (ReadPaused in transp.state)
+
 proc readExactly*(transp: StreamTransport, pbytes: pointer,
                  nbytes: int) {.async.} =
  ## Read exactly ``nbytes`` bytes from transport ``transp`` and store it to
@ -847,7 +859,7 @@ proc readExactly*(transp: StreamTransport, pbytes: pointer,
    if transp.offset == 0:
      if (ReadError in transp.state):
        raise transp.getError()
-      if (ReadEof in transp.state) or (ReadClosed in transp.state):
+      if ReadClosed in transp.state or transp.atEof():
        raise newException(TransportIncompleteError, "Data incomplete!")

    if transp.offset >= (nbytes - index):
@ -1028,13 +1040,11 @@ proc read*(transp: StreamTransport, n = -1): Future[seq[byte]] {.async.} =
  ## This procedure allocates buffer seq[byte] and return it as result.
  checkClosed(transp)
  checkPending(transp)
-
  result = newSeq[byte]()
-
  while true:
    if (ReadError in transp.state):
      raise transp.getError()
-    if {ReadEof, ReadClosed} * transp.state != {}:
+    if ReadClosed in transp.state or transp.atEof():
      break

    if transp.offset > 0:
@ -1049,11 +1059,10 @@ proc read*(transp: StreamTransport, n = -1): Future[seq[byte]] {.async.} =
      else:
        if transp.offset >= (n - s):
          # size of buffer data is more then we need, grabbing only part
-          let part = transp.offset - (n - s)
          result.setLen(n)
          copyMem(cast[pointer](addr result[s]), addr(transp.buffer[0]),
-                  part)
-          transp.shiftBuffer(part)
+                  n - s)
+          transp.shiftBuffer(n - s)
          break
        else:
          # there not enough data in buffer, grabbing all
@ -1070,10 +1079,42 @@ proc read*(transp: StreamTransport, n = -1): Future[seq[byte]] {.async.} =
    # Future[T], because readLoop continues working.
    transp.reader = nil

-proc atEof*(transp: StreamTransport): bool {.inline.} =
-  ## Returns ``true`` if ``transp`` is at EOF.
-  result = (transp.offset == 0) and (ReadEof in transp.state) and
-           (ReadPaused in transp.state)
+proc consume*(transp: StreamTransport, n = -1): Future[int] {.async.} =
+  ## Consume all bytes (n == -1) or ``n`` bytes from transport ``transp``.
+  ##
+  ## Return number of bytes actually consumed
+  checkClosed(transp)
+  checkPending(transp)
+  result = 0
+  while true:
+    if (ReadError in transp.state):
+      raise transp.getError()
+    if ReadClosed in transp.state or transp.atEof():
+      break
+    if transp.offset > 0:
+      if n == -1:
+        # consume all incoming data, until EOF
+        result += transp.offset
+        transp.offset = 0
+      else:
+        let left = n - result
+        if transp.offset >= left:
+          # size of buffer data is more then we need, consume only part
+          result += left
+          transp.shiftBuffer(left)
+          break
+        else:
+          # there not enough data in buffer, consume all
+          result += transp.offset
+          transp.offset = 0
+
+    transp.reader = newFuture[void]("stream.transport.consume")
+    if ReadPaused in transp.state:
+      transp.resumeRead()
+    await transp.reader
+    # we need to clear transp.reader to avoid double completion of this
+    # Future[T], because readLoop continues working.
+    transp.reader = nil

 proc join*(transp: StreamTransport) {.async.} =
  ## Wait until ``transp`` will not be closed.
--- a/tests/teststream.nim
+++ b/tests/teststream.nim
@ -15,6 +15,7 @@ else:
  import posix

 const
+  ConstantMessage = "SOMEDATA"
  ClientsCount = 100
  MessagesCount = 100
  MessageSize = 20
@ -100,6 +101,36 @@ proc serveClient4(server: StreamServer,
  var res = await transp.write(cast[pointer](addr answer[0]), len(answer))
  doAssert(res == len(answer))

+proc serveClient5(server: StreamServer,
+                  transp: StreamTransport, udata: pointer) {.async.} =
+  var data = await transp.read()
+  doAssert(len(data) == len(ConstantMessage) * MessagesCount)
+  transp.close()
+  var expect = ""
+  for i in 0..<MessagesCount:
+    expect.add(ConstantMessage)
+  doAssert(equalMem(addr expect[0], addr data[0], len(data)))
+  var counter = cast[ptr int](udata)
+  dec(counter[])
+  if counter[] == 0:
+    server.stop()
+    server.close()
+
+proc serveClient6(server: StreamServer,
+                  transp: StreamTransport, udata: pointer) {.async.} =
+  var expect = ConstantMessage
+  var skip = await transp.consume(len(ConstantMessage) * (MessagesCount - 1))
+  doAssert(skip == len(ConstantMessage) * (MessagesCount - 1))
+  var data = await transp.read()
+  doAssert(len(data) == len(ConstantMessage))
+  transp.close()
+  doAssert(equalMem(addr data[0], addr expect[0], len(expect)))
+  var counter = cast[ptr int](udata)
+  dec(counter[])
+  if counter[] == 0:
+    server.stop()
+    server.close()
+
 proc swarmWorker1(address: TransportAddress): Future[int] {.async.} =
  var transp = await connect(address)
  for i in 0..<MessagesCount:
@ -188,6 +219,24 @@ proc swarmWorker4(address: TransportAddress): Future[int] {.async.} =
  result = 1
  transp.close()

+proc swarmWorker5(address: TransportAddress): Future[int] {.async.} =
+  var transp = await connect(address)
+  var data = ConstantMessage
+  for i in 0..<MessagesCount:
+    var res = await transp.write(data)
+  result = MessagesCount
+  transp.close()
+
+proc swarmWorker6(address: TransportAddress): Future[int] {.async.} =
+  var transp = await connect(address)
+  var data = ConstantMessage
+  var seqdata = newSeq[byte](len(data))
+  copyMem(addr seqdata[0], addr data[0], len(data))
+  for i in 0..<MessagesCount:
+    var res = await transp.write(seqdata)
+  result = MessagesCount
+  transp.close()
+
 proc waitAll[T](futs: seq[Future[T]]): Future[void] =
  var counter = len(futs)
  var retFuture = newFuture[void]("waitAll")
@ -243,6 +292,28 @@ proc swarmManager4(address: TransportAddress): Future[int] {.async.} =
    var res = workers[i].read()
    result += res

+proc swarmManager5(address: TransportAddress): Future[int] {.async.} =
+  var retFuture = newFuture[void]("swarm.manager.read")
+  var workers = newSeq[Future[int]](ClientsCount)
+  var count = ClientsCount
+  for i in 0..<ClientsCount:
+    workers[i] = swarmWorker5(address)
+  await waitAll(workers)
+  for i in 0..<ClientsCount:
+    var res = workers[i].read()
+    result += res
+
+proc swarmManager6(address: TransportAddress): Future[int] {.async.} =
+  var retFuture = newFuture[void]("swarm.manager.consume")
+  var workers = newSeq[Future[int]](ClientsCount)
+  var count = ClientsCount
+  for i in 0..<ClientsCount:
+    workers[i] = swarmWorker6(address)
+  await waitAll(workers)
+  for i in 0..<ClientsCount:
+    var res = workers[i].read()
+    result += res
+
 proc test1(): Future[int] {.async.} =
  var ta = initTAddress("127.0.0.1:31344")
  var server = createStreamServer(ta, serveClient1, {ReuseAddr})
@ -271,13 +342,30 @@ proc test3(): Future[int] {.async.} =

 proc test4(): Future[int] {.async.} =
  var ta = initTAddress("127.0.0.1:31347")
-  var counter = 0
  var server = createStreamServer(ta, serveClient4, {ReuseAddr})
  server.start()
  result = await swarmManager4(ta)
  server.stop()
  server.close()

+proc test5(): Future[int] {.async.} =
+  var ta = initTAddress("127.0.0.1:31347")
+  var counter = ClientsCount
+  var server = createStreamServer(ta, serveClient5, {ReuseAddr},
+                                  udata = cast[pointer](addr counter))
+  server.start()
+  result = await swarmManager5(ta)
+  await server.join()
+
+proc test6(): Future[int] {.async.} =
+  var ta = initTAddress("127.0.0.1:31347")
+  var counter = ClientsCount
+  var server = createStreamServer(ta, serveClient6, {ReuseAddr},
+                                  udata = cast[pointer](addr counter))
+  server.start()
+  result = await swarmManager6(ta)
+  await server.join()
+
 when isMainModule:
  const
    m1 = "readLine() multiple clients with messages (" & $ClientsCount &
@ -287,6 +375,10 @@ when isMainModule:
    m3 = "readUntil() multiple clients with messages (" & $ClientsCount &
         " clients x " & $MessagesCount & " messages)"
    m4 = "writeFile() multiple clients (" & $FilesCount & " files)"
+    m5 = "write(string)/read(int) multiple clients (" & $ClientsCount &
+         " clients x " & $MessagesCount & " messages)"
+    m6 = "write(seq[byte])/consume(int)/read(int) multiple clients (" &
+         $ClientsCount & " clients x " & $MessagesCount & " messages)"

  suite "Stream Transport test suite":
    test m1:
@ -297,3 +389,7 @@ when isMainModule:
      check waitFor(test3()) == ClientsCount * MessagesCount
    test m4:
      check waitFor(test4()) == FilesCount
+    test m5:
+      check waitFor(test5()) == ClientsCount * MessagesCount
+    test m6:
+      check waitFor(test6()) == ClientsCount * MessagesCount