You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

287 lines
7.2 KiB

  1. package main
  2. import (
  3. "crypto/sha256"
  4. "encoding/binary"
  5. "encoding/hex"
  6. "encoding/json"
  7. "fmt"
  8. // it is ok to use math/rand here: we do not need a cryptographically secure random
  9. // number generator here and we can run the tests a bit faster
  10. "math/rand"
  11. "net"
  12. "net/http"
  13. "net/url"
  14. "os"
  15. "sync"
  16. "time"
  17. "github.com/gorilla/websocket"
  18. "github.com/pkg/errors"
  19. "github.com/tendermint/tendermint/libs/log"
  20. rpctypes "github.com/tendermint/tendermint/rpc/lib/types"
  21. )
  22. const (
  23. sendTimeout = 10 * time.Second
  24. // see https://github.com/tendermint/tendermint/blob/master/rpc/lib/server/handlers.go
  25. pingPeriod = (30 * 9 / 10) * time.Second
  26. )
  27. type transacter struct {
  28. Target string
  29. Rate int
  30. Size int
  31. Connections int
  32. BroadcastTxMethod string
  33. conns []*websocket.Conn
  34. connsBroken []bool
  35. startingWg sync.WaitGroup
  36. endingWg sync.WaitGroup
  37. stopped bool
  38. logger log.Logger
  39. }
  40. func newTransacter(target string, connections, rate int, size int, broadcastTxMethod string) *transacter {
  41. return &transacter{
  42. Target: target,
  43. Rate: rate,
  44. Size: size,
  45. Connections: connections,
  46. BroadcastTxMethod: broadcastTxMethod,
  47. conns: make([]*websocket.Conn, connections),
  48. connsBroken: make([]bool, connections),
  49. logger: log.NewNopLogger(),
  50. }
  51. }
  52. // SetLogger lets you set your own logger
  53. func (t *transacter) SetLogger(l log.Logger) {
  54. t.logger = l
  55. }
  56. // Start opens N = `t.Connections` connections to the target and creates read
  57. // and write goroutines for each connection.
  58. func (t *transacter) Start() error {
  59. t.stopped = false
  60. rand.Seed(time.Now().Unix())
  61. for i := 0; i < t.Connections; i++ {
  62. c, _, err := connect(t.Target) // nolint:bodyclose
  63. if err != nil {
  64. return err
  65. }
  66. t.conns[i] = c
  67. }
  68. t.startingWg.Add(t.Connections)
  69. t.endingWg.Add(2 * t.Connections)
  70. for i := 0; i < t.Connections; i++ {
  71. go t.sendLoop(i)
  72. go t.receiveLoop(i)
  73. }
  74. t.startingWg.Wait()
  75. return nil
  76. }
  77. // Stop closes the connections.
  78. func (t *transacter) Stop() {
  79. t.stopped = true
  80. t.endingWg.Wait()
  81. for _, c := range t.conns {
  82. c.Close()
  83. }
  84. }
  85. // receiveLoop reads messages from the connection (empty in case of
  86. // `broadcast_tx_async`).
  87. func (t *transacter) receiveLoop(connIndex int) {
  88. c := t.conns[connIndex]
  89. defer t.endingWg.Done()
  90. for {
  91. _, _, err := c.ReadMessage()
  92. if err != nil {
  93. if !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
  94. t.logger.Error(
  95. fmt.Sprintf("failed to read response on conn %d", connIndex),
  96. "err",
  97. err,
  98. )
  99. }
  100. return
  101. }
  102. if t.stopped || t.connsBroken[connIndex] {
  103. return
  104. }
  105. }
  106. }
  107. // sendLoop generates transactions at a given rate.
  108. func (t *transacter) sendLoop(connIndex int) {
  109. started := false
  110. // Close the starting waitgroup, in the event that this fails to start
  111. defer func() {
  112. if !started {
  113. t.startingWg.Done()
  114. }
  115. }()
  116. c := t.conns[connIndex]
  117. c.SetPingHandler(func(message string) error {
  118. err := c.WriteControl(websocket.PongMessage, []byte(message), time.Now().Add(sendTimeout))
  119. if err == websocket.ErrCloseSent {
  120. return nil
  121. } else if e, ok := err.(net.Error); ok && e.Temporary() {
  122. return nil
  123. }
  124. return err
  125. })
  126. logger := t.logger.With("addr", c.RemoteAddr())
  127. var txNumber = 0
  128. pingsTicker := time.NewTicker(pingPeriod)
  129. txsTicker := time.NewTicker(1 * time.Second)
  130. defer func() {
  131. pingsTicker.Stop()
  132. txsTicker.Stop()
  133. t.endingWg.Done()
  134. }()
  135. // hash of the host name is a part of each tx
  136. var hostnameHash [sha256.Size]byte
  137. hostname, err := os.Hostname()
  138. if err != nil {
  139. hostname = "127.0.0.1"
  140. }
  141. hostnameHash = sha256.Sum256([]byte(hostname))
  142. // each transaction embeds connection index, tx number and hash of the hostname
  143. // we update the tx number between successive txs
  144. tx := generateTx(connIndex, txNumber, t.Size, hostnameHash)
  145. txHex := make([]byte, len(tx)*2)
  146. hex.Encode(txHex, tx)
  147. for {
  148. select {
  149. case <-txsTicker.C:
  150. startTime := time.Now()
  151. endTime := startTime.Add(time.Second)
  152. numTxSent := t.Rate
  153. if !started {
  154. t.startingWg.Done()
  155. started = true
  156. }
  157. now := time.Now()
  158. for i := 0; i < t.Rate; i++ {
  159. // update tx number of the tx, and the corresponding hex
  160. updateTx(tx, txHex, txNumber)
  161. paramsJSON, err := json.Marshal(map[string]interface{}{"tx": txHex})
  162. if err != nil {
  163. fmt.Printf("failed to encode params: %v\n", err)
  164. os.Exit(1)
  165. }
  166. rawParamsJSON := json.RawMessage(paramsJSON)
  167. c.SetWriteDeadline(now.Add(sendTimeout))
  168. err = c.WriteJSON(rpctypes.RPCRequest{
  169. JSONRPC: "2.0",
  170. ID: rpctypes.JSONRPCStringID("tm-bench"),
  171. Method: t.BroadcastTxMethod,
  172. Params: rawParamsJSON,
  173. })
  174. if err != nil {
  175. err = errors.Wrap(err,
  176. fmt.Sprintf("txs send failed on connection #%d", connIndex))
  177. t.connsBroken[connIndex] = true
  178. logger.Error(err.Error())
  179. return
  180. }
  181. // cache the time.Now() reads to save time.
  182. if i%5 == 0 {
  183. now = time.Now()
  184. if now.After(endTime) {
  185. // Plus one accounts for sending this tx
  186. numTxSent = i + 1
  187. break
  188. }
  189. }
  190. txNumber++
  191. }
  192. timeToSend := time.Since(startTime)
  193. logger.Info(fmt.Sprintf("sent %d transactions", numTxSent), "took", timeToSend)
  194. if timeToSend < 1*time.Second {
  195. sleepTime := time.Second - timeToSend
  196. logger.Debug(fmt.Sprintf("connection #%d is sleeping for %f seconds", connIndex, sleepTime.Seconds()))
  197. time.Sleep(sleepTime)
  198. }
  199. case <-pingsTicker.C:
  200. // go-rpc server closes the connection in the absence of pings
  201. c.SetWriteDeadline(time.Now().Add(sendTimeout))
  202. if err := c.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
  203. err = errors.Wrap(err,
  204. fmt.Sprintf("failed to write ping message on conn #%d", connIndex))
  205. logger.Error(err.Error())
  206. t.connsBroken[connIndex] = true
  207. }
  208. }
  209. if t.stopped {
  210. // To cleanly close a connection, a client should send a close
  211. // frame and wait for the server to close the connection.
  212. c.SetWriteDeadline(time.Now().Add(sendTimeout))
  213. err := c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
  214. if err != nil {
  215. err = errors.Wrap(err,
  216. fmt.Sprintf("failed to write close message on conn #%d", connIndex))
  217. logger.Error(err.Error())
  218. t.connsBroken[connIndex] = true
  219. }
  220. return
  221. }
  222. }
  223. }
  224. func connect(host string) (*websocket.Conn, *http.Response, error) {
  225. u := url.URL{Scheme: "ws", Host: host, Path: "/websocket"}
  226. return websocket.DefaultDialer.Dial(u.String(), nil)
  227. }
  228. func generateTx(connIndex int, txNumber int, txSize int, hostnameHash [sha256.Size]byte) []byte {
  229. tx := make([]byte, txSize)
  230. binary.PutUvarint(tx[:8], uint64(connIndex))
  231. binary.PutUvarint(tx[8:16], uint64(txNumber))
  232. copy(tx[16:32], hostnameHash[:16])
  233. binary.PutUvarint(tx[32:40], uint64(time.Now().Unix()))
  234. // 40-* random data
  235. if _, err := rand.Read(tx[40:]); err != nil { //nolint: gosec
  236. panic(errors.Wrap(err, "failed to read random bytes"))
  237. }
  238. return tx
  239. }
  240. // warning, mutates input byte slice
  241. func updateTx(tx []byte, txHex []byte, txNumber int) {
  242. binary.PutUvarint(tx[8:16], uint64(txNumber))
  243. hexUpdate := make([]byte, 16)
  244. hex.Encode(hexUpdate, tx[8:16])
  245. for i := 16; i < 32; i++ {
  246. txHex[i] = hexUpdate[i-16]
  247. }
  248. }