test-plans/perf/impl/nim-libp2p/v0.1/perf.nim at 0ab8af3c4409dbb586f605bacf802c3527715efd · vladopajic/test-plans · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import os, strutils, strformat, json
import chronos, bearssl/[rand, hash]
import ./nimlibp2p/libp2p
import
  ./nimlibp2p/libp2p/[protocols/perf/client, protocols/perf/server, protocols/perf/core]

const fixedPeerId = "12D3KooWPnQpbXGqzgESFrkaFh1xvCrB64ADnLQQRYfMhnbSuFHF"

type Flags = object
  runServer: bool
  serverIpAddress: TransportAddress
  transport: string
  uploadBytes: uint
  downloadBytes: uint

proc initFlagsFromParams(flags: var Flags) =
  var i = 1
  while i < paramCount():
    case paramStr(i)
    of "--run-server":
      flags.runServer = true
    of "--server-ip-address":
      flags.serverIpAddress = initTAddress(paramStr(i + 1))
      i += 1
    of "--transport":
      flags.transport = paramStr(i + 1)
      i += 1
    of "--upload-bytes":
      flags.uploadBytes = parseUInt(paramStr(i + 1))
      i += 1
    of "--download-bytes":
      flags.downloadBytes = parseUInt(paramStr(i + 1))
      i += 1
    else:
      discard
    i += 1

proc seededRng(): ref HmacDrbgContext =
  var seed: cint = 0
  var rng = (ref HmacDrbgContext)()
  hmacDrbgInit(rng[], addr sha256Vtable, cast[pointer](addr seed), sizeof(seed).uint)
  return rng

proc runServer(f: Flags) {.async.} =
  let endlessFut = newFuture[void]()
  var switch = SwitchBuilder
    .new()
    .withRng(seededRng())
    .withAddresses(@[MultiAddress.init(f.serverIpAddress).tryGet()])
    .withTcpTransport()
    # .withQuicTransport()
    .withMplex()
    .withNoise()
    .build()
  switch.mount(Perf.new())
  await switch.start()
  await endlessFut # Await forever, exit on interrupt

proc runClient(f: Flags) {.async.} =
  let switchBuilder = SwitchBuilder
    .new()
    .withRng(newRng())
    .withAddress(MultiAddress.init("/ip4/127.0.0.1/tcp/0").tryGet())
    .withMplex()
    .withNoise()
  let switch =
    case f.transport
    of "tcp":
      switchBuilder.withTcpTransport().build()
    # of "quic-v1": switchBuilder.withQuicTransport().build()
    else:
      raise (ref Defect)()
  await switch.start()

  let startTime = Moment.now()
  let conn = await switch.dial(
    PeerId.init(fixedPeerId).tryGet(),
    @[MultiAddress.init(f.serverIpAddress).tryGet()],
    PerfCodec,
  )
  discard await PerfClient.perf(conn, f.uploadBytes, f.downloadBytes)

  let dur = Moment.now() - startTime
  let resultFinal =
    %*{
      "type": "final",
      "timeSeconds": dur.seconds,
      "uploadBytes": f.uploadBytes,
      "downloadBytes": f.downloadBytes,
    }
  echo $resultFinal

proc main() {.async.} =
  var flags = Flags()
  flags.initFlagsFromParams()

  if flags.runServer:
    await runServer(flags)
  else:
    await runClient(flags)

waitFor(main())