From e03547ea3e2f0372c540c586e993803299d3c4b6 Mon Sep 17 00:00:00 2001
From: Ludovic Chenut <ludovic@status.im>
Date: Mon, 14 Aug 2023 17:25:55 +0200
Subject: [PATCH] Perf protocol (#925)

---
 libp2p/protocols/perf/client.nim | 47 +++++++++++++++++++++++++
 libp2p/protocols/perf/core.nim   | 14 ++++++++
 libp2p/protocols/perf/server.nim | 60 ++++++++++++++++++++++++++++++++
 3 files changed, 121 insertions(+)
 create mode 100644 libp2p/protocols/perf/client.nim
 create mode 100644 libp2p/protocols/perf/core.nim
 create mode 100644 libp2p/protocols/perf/server.nim

diff --git a/libp2p/protocols/perf/client.nim b/libp2p/protocols/perf/client.nim
new file mode 100644
index 0000000..467aac4
--- /dev/null
+++ b/libp2p/protocols/perf/client.nim
@@ -0,0 +1,47 @@
+# Nim-LibP2P
+# Copyright (c) 2023 Status Research & Development GmbH
+# Licensed under either of
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
+# at your option.
+# This file may not be copied, modified, or distributed except according to
+# those terms.
+
+## `Perf <https://github.com/libp2p/specs/blob/master/perf/perf.md>`_ protocol specification
+
+import chronos, chronicles, sequtils
+import stew/endians2
+import ./core, ../../stream/connection
+
+logScope:
+  topics = "libp2p perf"
+
+type PerfClient* = ref object of RootObj
+
+proc perf*(_: typedesc[PerfClient], conn: Connection,
+           sizeToWrite: uint64 = 0, sizeToRead: uint64 = 0):
+           Future[Duration] {.async, public.} =
+  var
+    size = sizeToWrite
+    buf: array[PerfSize, byte]
+  let start = Moment.now()
+  trace "starting performance benchmark", conn, sizeToWrite, sizeToRead
+
+  await conn.write(toSeq(toBytesBE(sizeToRead)))
+  while size > 0:
+    let toWrite = min(size, PerfSize)
+    await conn.write(buf[0..<toWrite])
+    size -= toWrite
+
+  await conn.close()
+
+  size = sizeToRead
+
+  while size > 0:
+    let toRead = min(size, PerfSize)
+    await conn.readExactly(addr buf[0], toRead.int)
+    size = size - toRead
+
+  let duration = Moment.now() - start
+  trace "finishing performance benchmark", duration
+  return duration
diff --git a/libp2p/protocols/perf/core.nim b/libp2p/protocols/perf/core.nim
new file mode 100644
index 0000000..bb61965
--- /dev/null
+++ b/libp2p/protocols/perf/core.nim
@@ -0,0 +1,14 @@
+# Nim-LibP2P
+# Copyright (c) 2023 Status Research & Development GmbH
+# Licensed under either of
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
+# at your option.
+# This file may not be copied, modified, or distributed except according to
+# those terms.
+
+## `Perf <https://github.com/libp2p/specs/blob/master/perf/perf.md>`_ protocol specification
+
+const
+  PerfCodec* = "/perf/1.0.0"
+  PerfSize* = 65536
diff --git a/libp2p/protocols/perf/server.nim b/libp2p/protocols/perf/server.nim
new file mode 100644
index 0000000..383c323
--- /dev/null
+++ b/libp2p/protocols/perf/server.nim
@@ -0,0 +1,60 @@
+# Nim-LibP2P
+# Copyright (c) 2023 Status Research & Development GmbH
+# Licensed under either of
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
+# at your option.
+# This file may not be copied, modified, or distributed except according to
+# those terms.
+
+## `Perf <https://github.com/libp2p/specs/blob/master/perf/perf.md>`_ protocol specification
+
+{.push raises: [].}
+
+import chronos, chronicles
+import stew/endians2
+import ./core,
+       ../protocol,
+       ../../stream/connection,
+       ../../utility
+
+export chronicles, connection
+
+logScope:
+  topics = "libp2p perf"
+
+type Perf* = ref object of LPProtocol
+
+proc new*(T: typedesc[Perf]): T {.public.} =
+  var p = T()
+  proc handle(conn: Connection, proto: string) {.async, gcsafe, closure.} =
+    var bytesRead = 0
+    try:
+      trace "Received benchmark performance check", conn
+      var
+        sizeBuffer: array[8, byte]
+        size: uint64
+      await conn.readExactly(addr sizeBuffer[0], 8)
+      size = uint64.fromBytesBE(sizeBuffer)
+
+      var toReadBuffer: array[PerfSize, byte]
+      try:
+        while true:
+          bytesRead += await conn.readOnce(addr toReadBuffer[0], PerfSize)
+      except CatchableError as exc:
+        discard
+
+      var buf: array[PerfSize, byte]
+      while size > 0:
+        let toWrite = min(size, PerfSize)
+        await conn.write(buf[0..<toWrite])
+        size -= toWrite
+    except CancelledError as exc:
+      raise exc
+    except CatchableError as exc:
+      trace "exception in perf handler", exc = exc.msg, conn
+    await conn.close()
+
+  p.handler = handle
+  p.codec = PerfCodec
+  return p