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

Added consume()
Added tests for write(string), write(seq[byte]), consume().

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)
+                  n - s)
-          transp.shiftBuffer(part)
+          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.} =
+proc consume*(transp: StreamTransport, n = -1): Future[int] {.async.} =
-  ## Returns ``true`` if ``transp`` is at EOF.
+  ## Consume all bytes (n == -1) or ``n`` bytes from transport ``transp``.
-  result = (transp.offset == 0) and (ReadEof in transp.state) and
+  ##
-           (ReadPaused in transp.state)
+  ## 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.

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