Feature: Set weight on RPC methods

This commit makes weight on rpc method configurable through command line
argument `-m/--method`.
Default weights remain the same.

Example: `ethspam -m eth_blockNumber:1 -m eth_call:99` is going to send
1% of `eth_blockNumber` query and 99% of `eth_call`.

main.go

@@ -19,7 +19,8 @@ var Version = "dev"
 
 // Options contains the flag options
 type Options struct {
-	Web3Endpoint string `long:"rpc" description:"Ethereum JSONRPC provider, such as Infura or Cloudflare" default:"https://mainnet.infura.io/v3/af500e495f2d4e7cbcae36d0bfa66bcb"` // Versus API key on Infura
+	Methods      map[string]int64 `short:"m" long:"method" description:"A map from json rpc methods to their weight" default:"eth_getCode:100" default:"eth_getLogs:250" default:"eth_getTransactionByHash:250" default:"eth_blockNumber:350" default:"eth_getTransactionCount:400" default:"eth_getBlockByNumber:400" default:"eth_getBalance:550" default:"eth_getTransactionReceipt:600" default:"eth_call:2000"`
+	Web3Endpoint string           `long:"rpc" description:"Ethereum JSONRPC provider, such as Infura or Cloudflare" default:"https://mainnet.infura.io/v3/af500e495f2d4e7cbcae36d0bfa66bcb"` // Versus API key on Infura
 
 	Version bool `long:"version" description:"Print version and exit."`
 }
@@ -45,7 +46,10 @@ func main() {
 	}
 
 	gen := generator{}
-	installDefaults(&gen)
+	err = installDefaults(&gen, options.Methods)
+	if err != nil {
+		exit(1, "failed to install defaults: %s", err)
+	}
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

queries.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"strings"
@@ -79,7 +80,7 @@ func genEthGetCode(w io.Writer, s State) error {
 	return err
 }
 
-func installDefaults(gen *generator) {
+func installDefaults(gen *generator, methods map[string]int64) error {
 	// Top queries by weight, pulled from a 5000 Infura query sample on Dec 2019.
 	//     3 "eth_accounts"
 	//     4 "eth_getStorageAt"
@@ -102,57 +103,28 @@ func installDefaults(gen *generator) {
 	//   607 "eth_getTransactionReceipt"
 	//  1928 "eth_call"
 
-	gen.Add(RandomQuery{
-		Method:   "eth_call",
-		Weight:   2000,
-		Generate: genEthCall,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_getTransactionReceipt",
-		Weight:   600,
-		Generate: genEthGetTransactionReceipt,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_getBalance",
-		Weight:   550,
-		Generate: genEthGetBalance,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_getBlockByNumber",
-		Weight:   400,
-		Generate: genEthGetBalance,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_getTransactionCount",
-		Weight:   400,
-		Generate: genEthGetTransactionCount,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_blockNumber",
-		Weight:   350,
-		Generate: genEthBlockNumber,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_getTransactionByHash",
-		Weight:   250,
-		Generate: genEthGetTransactionByHash,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_getLogs",
-		Weight:   250,
-		Generate: genEthGetLogs,
-	})
-
-	gen.Add(RandomQuery{
-		Method:   "eth_getCode",
-		Weight:   100,
-		Generate: genEthGetCode,
-	})
+	rpcMethod := map[string]func(io.Writer, State) error{
+		"eth_call":                  genEthCall,
+		"eth_getTransactionReceipt": genEthGetTransactionReceipt,
+		"eth_getBalance":            genEthGetBalance,
+		"eth_getBlockByNumber":      genEthGetBlockByNumber,
+		"eth_getTransactionCount":   genEthGetTransactionCount,
+		"eth_blockNumber":           genEthBlockNumber,
+		"eth_getTransactionByHash":  genEthGetTransactionByHash,
+		"eth_getLogs":               genEthGetLogs,
+		"eth_getCode":               genEthGetCode,
+	}
+
+	for method, weight := range methods {
+		if _, err := rpcMethod[method]; err == false {
+			return errors.New(method + " is not supported")
+		}
+		gen.Add(RandomQuery{
+			Method:   method,
+			Weight:   weight,
+			Generate: rpcMethod[method],
+		})
+	}
+
+	return nil
 }