Merge #3641

3641: [Benchmark] Support multiple keys r=SaveTheRbtz a=SaveTheRbtz

This PR drastically speeds up account creation by using multiple keys on the service account.

Full list of changes from the `@peterargue's` original PR:
https://github.com/onflow/flow-go/compare/770759375d00e898798d38881153cdf0bf31fb03..rbtz/multiKey

Co-authored-by: Peter Argue <89119817+peterargue@users.noreply.github.com>
Co-authored-by: Alexey Ivanov <rbtz@dapperlabs.com>

integration/benchmark/account/account.go

@@ -0,0 +1,85 @@
+package account
+
+import (
+	"context"
+	"crypto/rand"
+	"fmt"
+
+	flowsdk "github.com/onflow/flow-go-sdk"
+
+	"github.com/onflow/flow-go-sdk/access"
+	"github.com/onflow/flow-go-sdk/crypto"
+)
+
+type FlowAccount struct {
+	Address *flowsdk.Address
+	ID      int
+
+	keys *keystore
+}
+
+func New(i int, address *flowsdk.Address, privKey crypto.PrivateKey, accountKeys []*flowsdk.AccountKey) (*FlowAccount, error) {
+	keys := make([]*accountKey, 0, len(accountKeys))
+	for _, key := range accountKeys {
+		signer, err := crypto.NewInMemorySigner(privKey, key.HashAlgo)
+		if err != nil {
+			return nil, fmt.Errorf("error while creating signer: %w", err)
+		}
+
+		keys = append(keys, &accountKey{
+			AccountKey: *key,
+			Address:    address,
+			Signer:     signer,
+		})
+	}
+
+	return &FlowAccount{
+		Address: address,
+		ID:      i,
+		keys:    newKeystore(keys),
+	}, nil
+}
+
+func LoadServiceAccount(
+	ctx context.Context,
+	flowClient access.Client,
+	servAccAddress *flowsdk.Address,
+	servAccPrivKeyHex string,
+) (*FlowAccount, error) {
+	acc, err := flowClient.GetAccount(ctx, *servAccAddress)
+	if err != nil {
+		return nil, fmt.Errorf("error while calling get account for service account: %w", err)
+	}
+
+	privateKey, err := crypto.DecodePrivateKeyHex(acc.Keys[0].SigAlgo, servAccPrivKeyHex)
+	if err != nil {
+		return nil, fmt.Errorf("error while decoding serice account private key hex: %w", err)
+	}
+
+	return New(0, servAccAddress, privateKey, acc.Keys)
+}
+
+func (acc *FlowAccount) NumKeys() int {
+	return acc.keys.Size()
+}
+
+func (acc *FlowAccount) GetKey() (*accountKey, error) {
+	return acc.keys.getKey()
+}
+
+// randomPrivateKey returns a randomly generated ECDSA P-256 private key.
+func RandomPrivateKey() crypto.PrivateKey {
+	seed := make([]byte, crypto.MinSeedLength)
+
+	_, err := rand.Read(seed)
+	if err != nil {
+		panic(err)
+	}
+
+	privateKey, err := crypto.GeneratePrivateKey(crypto.ECDSA_P256, seed)
+	if err != nil {
+		panic(err)
+	}
+
+	return privateKey
+}

integration/benchmark/account/keys.go

@@ -0,0 +1,102 @@
+package account
+
+import (
+	"fmt"
+	"sync"
+
+	flowsdk "github.com/onflow/flow-go-sdk"
+
+	"github.com/onflow/flow-go-sdk/crypto"
+)
+
+var ErrNoKeysAvailable = fmt.Errorf("no keys available")
+
+type accountKey struct {
+	flowsdk.AccountKey
+
+	mu      sync.Mutex
+	ks      *keystore
+	Address *flowsdk.Address
+	Signer  crypto.InMemorySigner
+	inuse   bool
+}
+
+type keystore struct {
+	availableKeys chan *accountKey
+	size          int
+}
+
+func newKeystore(keys []*accountKey) *keystore {
+	ks := &keystore{}
+
+	availableKeys := make(chan *accountKey, len(keys))
+	for _, key := range keys {
+		key.ks = ks
+		availableKeys <- key
+	}
+	ks.size = len(keys)
+	ks.availableKeys = availableKeys
+
+	return ks
+}
+
+func (k *keystore) Size() int {
+	return k.size
+}
+
+func (k *keystore) getKey() (*accountKey, error) {
+	select {
+	case key := <-k.availableKeys:
+		key.mu.Lock()
+		defer key.mu.Unlock()
+
+		if key.inuse {
+			return nil, fmt.Errorf("key already in use")
+		}
+		key.inuse = true
+
+		return key, nil
+	default:
+		return nil, ErrNoKeysAvailable
+	}
+}
+
+func (k *accountKey) markUnused() {
+	k.mu.Lock()
+	defer k.mu.Unlock()
+
+	k.inuse = false
+}
+
+// Done unlocks a key after use and puts it back into the pool.
+func (k *accountKey) Done() {
+	k.markUnused()
+	k.ks.availableKeys <- k
+}
+
+// IncrementSequenceNumber is called when a key was successfully used to sign a transaction as the proposer.
+// It increments the sequence number.
+func (k *accountKey) IncrementSequenceNumber() {
+	k.mu.Lock()
+	defer k.mu.Unlock()
+
+	if !k.inuse {
+		panic("assert: key not locked")
+	}
+	k.SequenceNumber++
+}
+
+func (k *accountKey) SignPayload(tx *flowsdk.Transaction) error {
+	return tx.SignPayload(*k.Address, k.Index, k.Signer)
+}
+
+func (k *accountKey) SignTx(tx *flowsdk.Transaction) error {
+	if len(tx.Authorizers) == 0 {
+		tx = tx.AddAuthorizer(*k.Address)
+	}
+
+	return tx.
+		SetProposalKey(*k.Address, k.Index, k.SequenceNumber).
+		SetPayer(*k.Address).
+		SignEnvelope(*k.Address, k.Index, k.Signer)
+}

integration/benchmark/cmd/ci/main.go

@@ -119,7 +119,14 @@ func main() {
 		Stringer("flowTokenAddress", flowTokenAddress).
 		Msg("addresses")
 
-	flowClient, err := client.NewClient(accessNodeAddress, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	flowClient, err := client.NewClient(
+		accessNodeAddress,
+		grpc.WithDefaultCallOptions(
+			grpc.MaxCallRecvMsgSize(defaultMaxMsgSize),
+			grpc.MaxCallSendMsgSize(defaultMaxMsgSize),
+		),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
 	if err != nil {
 		log.Fatal().Err(err).Msg("unable to initialize Flow client")
 	}

integration/benchmark/cmd/manual/main.go

@@ -29,6 +29,10 @@ type LoadCase struct {
 	duration time.Duration
 }
 
+const (
+	defaultMaxMsgSize = 1024 * 1024 * 16 // 16 MB
+)
+
 func main() {
 	sleep := flag.Duration("sleep", 0, "duration to sleep before benchmarking starts")
 	loadTypeFlag := flag.String("load-type", "token-transfer", "type of loads (\"token-transfer\", \"add-keys\", \"computation-heavy\", \"event-heavy\", \"ledger-heavy\", \"const-exec\")")
@@ -86,7 +90,14 @@ func main() {
 	accessNodeAddrs := strings.Split(*accessNodes, ",")
 	clients := make([]access.Client, 0, len(accessNodeAddrs))
 	for _, addr := range accessNodeAddrs {
-		client, err := client.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+		client, err := client.NewClient(
+			addr,
+			grpc.WithDefaultCallOptions(
+				grpc.MaxCallRecvMsgSize(defaultMaxMsgSize),
+				grpc.MaxCallSendMsgSize(defaultMaxMsgSize),
+			),
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+		)
 		if err != nil {
 			log.Fatal().Str("addr", addr).Err(err).Msgf("unable to initialize flow client")
 		}