Merge pull request #550 from nspcc-dev/cache-publickeys

vm/core: improve block import speed with PublicKey caching

pkg/core/blockchain.go

@@ -73,6 +73,9 @@ type Blockchain struct {
 	runToExitCh chan struct{}
 
 	memPool MemPool
+
+	// cache for block verification keys.
+	keyCache map[util.Uint160]map[string]*keys.PublicKey
 }
 
 type headersOpFunc func(headerList *HeaderHashList)
@@ -88,6 +91,7 @@ func NewBlockchain(s storage.Store, cfg config.ProtocolConfiguration) (*Blockcha
 		stopCh:        make(chan struct{}),
 		runToExitCh:   make(chan struct{}),
 		memPool:       NewMemPool(50000),
+		keyCache:      make(map[util.Uint160]map[string]*keys.PublicKey),
 	}
 
 	if err := bc.init(); err != nil {
@@ -1427,7 +1431,7 @@ func (bc *Blockchain) GetTestVM() (*vm.VM, storage.Store) {
 }
 
 // verifyHashAgainstScript verifies given hash against the given witness.
-func (bc *Blockchain) verifyHashAgainstScript(hash util.Uint160, witness *transaction.Witness, checkedHash util.Uint256, interopCtx *interopContext) error {
+func (bc *Blockchain) verifyHashAgainstScript(hash util.Uint160, witness *transaction.Witness, checkedHash util.Uint256, interopCtx *interopContext, useKeys bool) error {
 	verification := witness.VerificationScript
 
 	if len(verification) == 0 {
@@ -1447,6 +1451,9 @@ func (bc *Blockchain) verifyHashAgainstScript(hash util.Uint160, witness *transa
 	vm.SetCheckedHash(checkedHash.BytesBE())
 	vm.LoadScript(verification)
 	vm.LoadScript(witness.InvocationScript)
+	if useKeys && bc.keyCache[hash] != nil {
+		vm.SetPublicKeys(bc.keyCache[hash])
+	}
 	err := vm.Run()
 	if vm.HasFailed() {
 		return errors.Errorf("vm failed to execute the script with error: %s", err)
@@ -1460,6 +1467,9 @@ func (bc *Blockchain) verifyHashAgainstScript(hash util.Uint160, witness *transa
 		if !res {
 			return errors.Errorf("signature check failed")
 		}
+		if useKeys && bc.keyCache[hash] == nil {
+			bc.keyCache[hash] = vm.GetPublicKeys()
+		}
 	} else {
 		return errors.Errorf("no result returned from the script")
 	}
@@ -1487,7 +1497,7 @@ func (bc *Blockchain) verifyTxWitnesses(t *transaction.Transaction, block *Block
 	sort.Slice(witnesses, func(i, j int) bool { return witnesses[i].ScriptHash().Less(witnesses[j].ScriptHash()) })
 	interopCtx := newInteropContext(trigger.Verification, bc, bc.store, block, t)
 	for i := 0; i < len(hashes); i++ {
-		err := bc.verifyHashAgainstScript(hashes[i], &witnesses[i], t.VerificationHash(), interopCtx)
+		err := bc.verifyHashAgainstScript(hashes[i], &witnesses[i], t.VerificationHash(), interopCtx, false)
 		if err != nil {
 			numStr := fmt.Sprintf("witness #%d", i)
 			return errors.Wrap(err, numStr)
@@ -1506,7 +1516,7 @@ func (bc *Blockchain) verifyBlockWitnesses(block *Block, prevHeader *Header) err
 		hash = prevHeader.NextConsensus
 	}
 	interopCtx := newInteropContext(trigger.Verification, bc, bc.store, nil, nil)
-	return bc.verifyHashAgainstScript(hash, &block.Script, block.VerificationHash(), interopCtx)
+	return bc.verifyHashAgainstScript(hash, &block.Script, block.VerificationHash(), interopCtx, true)
 }
 
 func hashAndIndexToBytes(h util.Uint256, index uint32) []byte {

pkg/vm/vm.go

@@ -75,6 +75,9 @@ type VM struct {
 
 	itemCount map[StackItem]int
 	size      int
+
+	// Public keys cache.
+	keys map[string]*keys.PublicKey
 }
 
 // InteropFuncPrice represents an interop function with a price.
@@ -92,6 +95,7 @@ func New() *VM {
 		istack:    NewStack("invocation"),
 
 		itemCount: make(map[StackItem]int),
+		keys:      make(map[string]*keys.PublicKey),
 	}
 
 	vm.estack = vm.newItemStack("evaluation")
@@ -145,6 +149,17 @@ func (v *VM) Istack() *Stack {
 	return v.istack
 }
 
+// SetPublicKeys sets internal key cache to the specified value (note
+// that it doesn't copy them).
+func (v *VM) SetPublicKeys(keys map[string]*keys.PublicKey) {
+	v.keys = keys
+}
+
+// GetPublicKeys returns internal key cache (note that it doesn't copy it).
+func (v *VM) GetPublicKeys() map[string]*keys.PublicKey {
+	return v.keys
+}
+
 // LoadArgs loads in the arguments used in the Mian entry point.
 func (v *VM) LoadArgs(method []byte, args []StackItem) {
 	if len(args) > 0 {
@@ -1168,11 +1183,7 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 			msg := v.estack.Pop().Bytes()
 			hashToCheck = hash.Sha256(msg).BytesBE()
 		}
-		pkey := &keys.PublicKey{}
-		err := pkey.DecodeBytes(keyb)
-		if err != nil {
-			panic(err.Error())
-		}
+		pkey := v.bytesToPublicKey(keyb)
 		res := pkey.Verify(signature, hashToCheck)
 		v.estack.PushVal(res)
 
@@ -1197,11 +1208,7 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 		// j counts keys and i counts signatures.
 		j := 0
 		for i := 0; sigok && j < len(pkeys) && i < len(sigs); {
-			pkey := &keys.PublicKey{}
-			err := pkey.DecodeBytes(pkeys[j])
-			if err != nil {
-				panic(err.Error())
-			}
+			pkey := v.bytesToPublicKey(pkeys[j])
 			// We only move to the next signature if the check was
 			// successful, but if it's not maybe the next key will
 			// fit, so we always move to the next key.
@@ -1424,3 +1431,21 @@ func (v *VM) checkBigIntSize(a *big.Int) {
 		panic("big integer is too big")
 	}
 }
+
+// bytesToPublicKey is a helper deserializing keys using cache and panicing on
+// error.
+func (v *VM) bytesToPublicKey(b []byte) *keys.PublicKey {
+	var pkey *keys.PublicKey
+	s := string(b)
+	if v.keys[s] != nil {
+		pkey = v.keys[s]
+	} else {
+		pkey = &keys.PublicKey{}
+		err := pkey.DecodeBytes(b)
+		if err != nil {
+			panic(err.Error())
+		}
+		v.keys[s] = pkey
+	}
+	return pkey
+}

pkg/vm/vm_test.go

@@ -41,6 +41,25 @@ func TestRegisterInterop(t *testing.T) {
 	assert.Equal(t, true, ok)
 }
 
+func TestBytesToPublicKey(t *testing.T) {
+	v := New()
+	cache := v.GetPublicKeys()
+	assert.Equal(t, 0, len(cache))
+	keyHex := "03b209fd4f53a7170ea4444e0cb0a6bb6a53c2bd016926989cf85f9b0fba17a70c"
+	keyBytes, _ := hex.DecodeString(keyHex)
+	key := v.bytesToPublicKey(keyBytes)
+	assert.NotNil(t, key)
+	key2 := v.bytesToPublicKey(keyBytes)
+	assert.Equal(t, key, key2)
+
+	cache = v.GetPublicKeys()
+	assert.Equal(t, 1, len(cache))
+	assert.NotNil(t, cache[string(keyBytes)])
+
+	keyBytes[0] = 0xff
+	require.Panics(t, func() { v.bytesToPublicKey(keyBytes) })
+}
+
 func TestPushBytes1to75(t *testing.T) {
 	buf := new(bytes.Buffer)
 	for i := 1; i <= 75; i++ {