[NOD-1543] Optimize the performance of Count() in BlockHeaderStore and BlockStore (#1109)

* [NOD-1543] Optimize Count() in BlockHeaderStore.

* [NOD-1543] Optimize Count() in BlockStore.

* [NOD-1543] Fix commitCount.

* [NOD-1543] Explicitly initialize count to 0.

domain/consensus/database/serialization/dbobjects.pb.go

@@ -1629,6 +1629,100 @@ func (x *DbVirtualDiffParents) GetVirtualDiffParents() []*DbHash {
 	return nil
 }
 
+type DbBlockCount struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	Count uint64 `protobuf:"varint,1,opt,name=count,proto3" json:"count,omitempty"`
+}
+
+func (x *DbBlockCount) Reset() {
+	*x = DbBlockCount{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_dbobjects_proto_msgTypes[28]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *DbBlockCount) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*DbBlockCount) ProtoMessage() {}
+
+func (x *DbBlockCount) ProtoReflect() protoreflect.Message {
+	mi := &file_dbobjects_proto_msgTypes[28]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use DbBlockCount.ProtoReflect.Descriptor instead.
+func (*DbBlockCount) Descriptor() ([]byte, []int) {
+	return file_dbobjects_proto_rawDescGZIP(), []int{28}
+}
+
+func (x *DbBlockCount) GetCount() uint64 {
+	if x != nil {
+		return x.Count
+	}
+	return 0
+}
+
+type DbBlockHeaderCount struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	Count uint64 `protobuf:"varint,1,opt,name=count,proto3" json:"count,omitempty"`
+}
+
+func (x *DbBlockHeaderCount) Reset() {
+	*x = DbBlockHeaderCount{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_dbobjects_proto_msgTypes[29]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *DbBlockHeaderCount) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*DbBlockHeaderCount) ProtoMessage() {}
+
+func (x *DbBlockHeaderCount) ProtoReflect() protoreflect.Message {
+	mi := &file_dbobjects_proto_msgTypes[29]
+	if protoimpl.UnsafeEnabled && x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use DbBlockHeaderCount.ProtoReflect.Descriptor instead.
+func (*DbBlockHeaderCount) Descriptor() ([]byte, []int) {
+	return file_dbobjects_proto_rawDescGZIP(), []int{29}
+}
+
+func (x *DbBlockHeaderCount) GetCount() uint64 {
+	if x != nil {
+		return x.Count
+	}
+	return 0
+}
+
 var File_dbobjects_proto protoreflect.FileDescriptor
 
 var file_dbobjects_proto_rawDesc = []byte{
@@ -1858,10 +1952,15 @@ var file_dbobjects_proto_rawDesc = []byte{
 	0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x15, 0x2e, 0x73, 0x65, 0x72, 0x69, 0x61, 0x6c,
 	0x69, 0x7a, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x44, 0x62, 0x48, 0x61, 0x73, 0x68, 0x52, 0x12,
 	0x76, 0x69, 0x72, 0x74, 0x75, 0x61, 0x6c, 0x44, 0x69, 0x66, 0x66, 0x50, 0x61, 0x72, 0x65, 0x6e,
-	0x74, 0x73, 0x42, 0x2a, 0x5a, 0x28, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d,
+	0x74, 0x73, 0x22, 0x24, 0x0a, 0x0c, 0x44, 0x62, 0x42, 0x6c, 0x6f, 0x63, 0x6b, 0x43, 0x6f, 0x75,
-	0x2f, 0x6b, 0x61, 0x73, 0x70, 0x61, 0x6e, 0x65, 0x74, 0x2f, 0x6b, 0x61, 0x73, 0x70, 0x61, 0x64,
+	0x6e, 0x74, 0x12, 0x14, 0x0a, 0x05, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x2f, 0x73, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x69, 0x7a, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x62, 0x06,
+	0x04, 0x52, 0x05, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x22, 0x2a, 0x0a, 0x12, 0x44, 0x62, 0x42, 0x6c,
-	0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+	0x6f, 0x63, 0x6b, 0x48, 0x65, 0x61, 0x64, 0x65, 0x72, 0x43, 0x6f, 0x75, 0x6e, 0x74, 0x12, 0x14,
+	0x0a, 0x05, 0x63, 0x6f, 0x75, 0x6e, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x04, 0x52, 0x05, 0x63,
+	0x6f, 0x75, 0x6e, 0x74, 0x42, 0x2a, 0x5a, 0x28, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63,
+	0x6f, 0x6d, 0x2f, 0x6b, 0x61, 0x73, 0x70, 0x61, 0x6e, 0x65, 0x74, 0x2f, 0x6b, 0x61, 0x73, 0x70,
+	0x61, 0x64, 0x2f, 0x73, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x69, 0x7a, 0x61, 0x74, 0x69, 0x6f, 0x6e,
+	0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
 }
 
 var (
@@ -1876,7 +1975,7 @@ func file_dbobjects_proto_rawDescGZIP() []byte {
 	return file_dbobjects_proto_rawDescData
 }
 
-var file_dbobjects_proto_msgTypes = make([]protoimpl.MessageInfo, 28)
+var file_dbobjects_proto_msgTypes = make([]protoimpl.MessageInfo, 30)
 var file_dbobjects_proto_goTypes = []interface{}{
 	(*DbBlock)(nil),                     // 0: serialization.DbBlock
 	(*DbBlockHeader)(nil),               // 1: serialization.DbBlockHeader
@@ -1906,6 +2005,8 @@ var file_dbobjects_proto_goTypes = []interface{}{
 	(*DbHeaderTips)(nil),                // 25: serialization.DbHeaderTips
 	(*DbTips)(nil),                      // 26: serialization.DbTips
 	(*DbVirtualDiffParents)(nil),        // 27: serialization.DbVirtualDiffParents
+	(*DbBlockCount)(nil),                // 28: serialization.DbBlockCount
+	(*DbBlockHeaderCount)(nil),          // 29: serialization.DbBlockHeaderCount
 }
 var file_dbobjects_proto_depIdxs = []int32{
 	1,  // 0: serialization.DbBlock.header:type_name -> serialization.DbBlockHeader
@@ -2292,6 +2393,30 @@ func file_dbobjects_proto_init() {
 				return nil
 			}
 		}
+		file_dbobjects_proto_msgTypes[28].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*DbBlockCount); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_dbobjects_proto_msgTypes[29].Exporter = func(v interface{}, i int) interface{} {
+			switch v := v.(*DbBlockHeaderCount); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
 	}
 	type x struct{}
 	out := protoimpl.TypeBuilder{
@@ -2299,7 +2424,7 @@ func file_dbobjects_proto_init() {
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
 			RawDescriptor: file_dbobjects_proto_rawDesc,
 			NumEnums:      0,
-			NumMessages:   28,
+			NumMessages:   30,
 			NumExtensions: 0,
 			NumServices:   0,
 		},

domain/consensus/database/serialization/dbobjects.proto

@@ -150,3 +150,11 @@ message DbTips {
 message DbVirtualDiffParents {
   repeated DbHash virtualDiffParents = 1;
 }
+
+message DbBlockCount {
+  uint64 count = 1;
+}
+
+message DbBlockHeaderCount {
+  uint64 count = 1;
+}

domain/consensus/datastructures/blockheaderstore/blockheaderstore.go

@@ -9,85 +9,119 @@ import (
 )
 
 var bucket = dbkeys.MakeBucket([]byte("block-headers"))
+var countKey = dbkeys.MakeBucket().Key([]byte("block-headers-count"))
 
 // blockHeaderStore represents a store of blocks
 type blockHeaderStore struct {
 	staging  map[externalapi.DomainHash]*externalapi.DomainBlockHeader
 	toDelete map[externalapi.DomainHash]struct{}
+	count    uint64
 }
 
 // New instantiates a new BlockHeaderStore
-func New() model.BlockHeaderStore {
+func New(dbContext model.DBReader) (model.BlockHeaderStore, error) {
-	return &blockHeaderStore{
+	blockHeaderStore := &blockHeaderStore{
 		staging:  make(map[externalapi.DomainHash]*externalapi.DomainBlockHeader),
 		toDelete: make(map[externalapi.DomainHash]struct{}),
 	}
-}
 
-// Stage stages the given block header for the given blockHash
+	err := blockHeaderStore.initializeCount(dbContext)
-func (bms *blockHeaderStore) Stage(blockHash *externalapi.DomainHash, blockHeader *externalapi.DomainBlockHeader) error {
-	clone, err := bms.cloneHeader(blockHeader)
-	if err != nil {
-		return err
-	}
-
-	bms.staging[*blockHash] = clone
-	return nil
-}
-
-func (bms *blockHeaderStore) IsStaged() bool {
-	return len(bms.staging) != 0 || len(bms.toDelete) != 0
-}
-
-func (bms *blockHeaderStore) Discard() {
-	bms.staging = make(map[externalapi.DomainHash]*externalapi.DomainBlockHeader)
-	bms.toDelete = make(map[externalapi.DomainHash]struct{})
-}
-
-func (bms *blockHeaderStore) Commit(dbTx model.DBTransaction) error {
-	for hash, header := range bms.staging {
-		headerBytes, err := bms.serializeHeader(header)
-		if err != nil {
-			return err
-		}
-		err = dbTx.Put(bms.hashAsKey(&hash), headerBytes)
-		if err != nil {
-			return err
-		}
-	}
-
-	for hash := range bms.toDelete {
-		err := dbTx.Delete(bms.hashAsKey(&hash))
-		if err != nil {
-			return err
-		}
-	}
-
-	bms.Discard()
-	return nil
-}
-
-// BlockHeader gets the block header associated with the given blockHash
-func (bms *blockHeaderStore) BlockHeader(dbContext model.DBReader, blockHash *externalapi.DomainHash) (*externalapi.DomainBlockHeader, error) {
-	if header, ok := bms.staging[*blockHash]; ok {
-		return header, nil
-	}
-
-	headerBytes, err := dbContext.Get(bms.hashAsKey(blockHash))
 	if err != nil {
 		return nil, err
 	}
 
-	return bms.deserializeHeader(headerBytes)
+	return blockHeaderStore, nil
+}
+
+func (bhs *blockHeaderStore) initializeCount(dbContext model.DBReader) error {
+	count := uint64(0)
+	hasCountBytes, err := dbContext.Has(countKey)
+	if err != nil {
+		return err
+	}
+	if hasCountBytes {
+		countBytes, err := dbContext.Get(countKey)
+		if err != nil {
+			return err
+		}
+		count, err = bhs.deserializeHeaderCount(countBytes)
+		if err != nil {
+			return err
+		}
+	}
+	bhs.count = count
+	return nil
+}
+
+// Stage stages the given block header for the given blockHash
+func (bhs *blockHeaderStore) Stage(blockHash *externalapi.DomainHash, blockHeader *externalapi.DomainBlockHeader) error {
+	clone, err := bhs.cloneHeader(blockHeader)
+	if err != nil {
+		return err
+	}
+
+	bhs.staging[*blockHash] = clone
+	return nil
+}
+
+func (bhs *blockHeaderStore) IsStaged() bool {
+	return len(bhs.staging) != 0 || len(bhs.toDelete) != 0
+}
+
+func (bhs *blockHeaderStore) Discard() {
+	bhs.staging = make(map[externalapi.DomainHash]*externalapi.DomainBlockHeader)
+	bhs.toDelete = make(map[externalapi.DomainHash]struct{})
+}
+
+func (bhs *blockHeaderStore) Commit(dbTx model.DBTransaction) error {
+	for hash, header := range bhs.staging {
+		headerBytes, err := bhs.serializeHeader(header)
+		if err != nil {
+			return err
+		}
+		err = dbTx.Put(bhs.hashAsKey(&hash), headerBytes)
+		if err != nil {
+			return err
+		}
+	}
+
+	for hash := range bhs.toDelete {
+		err := dbTx.Delete(bhs.hashAsKey(&hash))
+		if err != nil {
+			return err
+		}
+	}
+
+	err := bhs.commitCount(dbTx)
+	if err != nil {
+		return err
+	}
+
+	bhs.Discard()
+	return nil
+}
+
+// BlockHeader gets the block header associated with the given blockHash
+func (bhs *blockHeaderStore) BlockHeader(dbContext model.DBReader, blockHash *externalapi.DomainHash) (*externalapi.DomainBlockHeader, error) {
+	if header, ok := bhs.staging[*blockHash]; ok {
+		return header, nil
+	}
+
+	headerBytes, err := dbContext.Get(bhs.hashAsKey(blockHash))
+	if err != nil {
+		return nil, err
+	}
+
+	return bhs.deserializeHeader(headerBytes)
 }
 
 // HasBlock returns whether a block header with a given hash exists in the store.
-func (bms *blockHeaderStore) HasBlockHeader(dbContext model.DBReader, blockHash *externalapi.DomainHash) (bool, error) {
+func (bhs *blockHeaderStore) HasBlockHeader(dbContext model.DBReader, blockHash *externalapi.DomainHash) (bool, error) {
-	if _, ok := bms.staging[*blockHash]; ok {
+	if _, ok := bhs.staging[*blockHash]; ok {
 		return true, nil
 	}
 
-	exists, err := dbContext.Has(bms.hashAsKey(blockHash))
+	exists, err := dbContext.Has(bhs.hashAsKey(blockHash))
 	if err != nil {
 		return false, err
 	}
@@ -96,11 +130,11 @@ func (bms *blockHeaderStore) HasBlockHeader(dbContext model.DBReader, blockHash
 }
 
 // BlockHeaders gets the block headers associated with the given blockHashes
-func (bms *blockHeaderStore) BlockHeaders(dbContext model.DBReader, blockHashes []*externalapi.DomainHash) ([]*externalapi.DomainBlockHeader, error) {
+func (bhs *blockHeaderStore) BlockHeaders(dbContext model.DBReader, blockHashes []*externalapi.DomainHash) ([]*externalapi.DomainBlockHeader, error) {
 	headers := make([]*externalapi.DomainBlockHeader, len(blockHashes))
 	for i, hash := range blockHashes {
 		var err error
-		headers[i], err = bms.BlockHeader(dbContext, hash)
+		headers[i], err = bhs.BlockHeader(dbContext, hash)
 		if err != nil {
 			return nil, err
 		}
@@ -109,24 +143,24 @@ func (bms *blockHeaderStore) BlockHeaders(dbContext model.DBReader, blockHashes
 }
 
 // Delete deletes the block associated with the given blockHash
-func (bms *blockHeaderStore) Delete(blockHash *externalapi.DomainHash) {
+func (bhs *blockHeaderStore) Delete(blockHash *externalapi.DomainHash) {
-	if _, ok := bms.staging[*blockHash]; ok {
+	if _, ok := bhs.staging[*blockHash]; ok {
-		delete(bms.staging, *blockHash)
+		delete(bhs.staging, *blockHash)
 		return
 	}
-	bms.toDelete[*blockHash] = struct{}{}
+	bhs.toDelete[*blockHash] = struct{}{}
 }
 
-func (bms *blockHeaderStore) hashAsKey(hash *externalapi.DomainHash) model.DBKey {
+func (bhs *blockHeaderStore) hashAsKey(hash *externalapi.DomainHash) model.DBKey {
 	return bucket.Key(hash[:])
 }
 
-func (bms *blockHeaderStore) serializeHeader(header *externalapi.DomainBlockHeader) ([]byte, error) {
+func (bhs *blockHeaderStore) serializeHeader(header *externalapi.DomainBlockHeader) ([]byte, error) {
 	dbBlockHeader := serialization.DomainBlockHeaderToDbBlockHeader(header)
 	return proto.Marshal(dbBlockHeader)
 }
 
-func (bms *blockHeaderStore) deserializeHeader(headerBytes []byte) (*externalapi.DomainBlockHeader, error) {
+func (bhs *blockHeaderStore) deserializeHeader(headerBytes []byte) (*externalapi.DomainBlockHeader, error) {
 	dbBlockHeader := &serialization.DbBlockHeader{}
 	err := proto.Unmarshal(headerBytes, dbBlockHeader)
 	if err != nil {
@@ -135,23 +169,43 @@ func (bms *blockHeaderStore) deserializeHeader(headerBytes []byte) (*externalapi
 	return serialization.DbBlockHeaderToDomainBlockHeader(dbBlockHeader)
 }
 
-func (bms *blockHeaderStore) cloneHeader(header *externalapi.DomainBlockHeader) (*externalapi.DomainBlockHeader, error) {
+func (bhs *blockHeaderStore) cloneHeader(header *externalapi.DomainBlockHeader) (*externalapi.DomainBlockHeader, error) {
-	serialized, err := bms.serializeHeader(header)
+	serialized, err := bhs.serializeHeader(header)
 	if err != nil {
 		return nil, err
 	}
 
-	return bms.deserializeHeader(serialized)
+	return bhs.deserializeHeader(serialized)
 }
 
-func (bms *blockHeaderStore) Count(dbContext model.DBReader) (uint64, error) {
+func (bhs *blockHeaderStore) Count() uint64 {
-	cursor, err := dbContext.Cursor(bucket)
+	return bhs.count + uint64(len(bhs.staging)) - uint64(len(bhs.toDelete))
+}
+
+func (bhs *blockHeaderStore) deserializeHeaderCount(countBytes []byte) (uint64, error) {
+	dbBlockHeaderCount := &serialization.DbBlockHeaderCount{}
+	err := proto.Unmarshal(countBytes, dbBlockHeaderCount)
 	if err != nil {
 		return 0, err
 	}
-	count := uint64(0)
+	return dbBlockHeaderCount.Count, nil
-	for cursor.Next() {
+}
-		count++
+
-	}
+func (bhs *blockHeaderStore) commitCount(dbTx model.DBTransaction) error {
-	return count, nil
+	count := bhs.Count()
+	countBytes, err := bhs.serializeHeaderCount(count)
+	if err != nil {
+		return err
+	}
+	err = dbTx.Put(countKey, countBytes)
+	if err != nil {
+		return err
+	}
+	bhs.count = count
+	return nil
+}
+
+func (bhs *blockHeaderStore) serializeHeaderCount(count uint64) ([]byte, error) {
+	dbBlockHeaderCount := &serialization.DbBlockHeaderCount{Count: count}
+	return proto.Marshal(dbBlockHeaderCount)
 }

domain/consensus/datastructures/blockstore/blockstore.go

@@ -9,85 +9,119 @@ import (
 )
 
 var bucket = dbkeys.MakeBucket([]byte("blocks"))
+var countKey = dbkeys.MakeBucket().Key([]byte("blocks-count"))
 
 // blockStore represents a store of blocks
 type blockStore struct {
 	staging  map[externalapi.DomainHash]*externalapi.DomainBlock
 	toDelete map[externalapi.DomainHash]struct{}
+	count    uint64
 }
 
 // New instantiates a new BlockStore
-func New() model.BlockStore {
+func New(dbContext model.DBReader) (model.BlockStore, error) {
-	return &blockStore{
+	blockStore := &blockStore{
 		staging:  make(map[externalapi.DomainHash]*externalapi.DomainBlock),
 		toDelete: make(map[externalapi.DomainHash]struct{}),
 	}
-}
 
-// Stage stages the given block for the given blockHash
+	err := blockStore.initializeCount(dbContext)
-func (bms *blockStore) Stage(blockHash *externalapi.DomainHash, block *externalapi.DomainBlock) error {
-	clone, err := bms.clone(block)
-	if err != nil {
-		return err
-	}
-
-	bms.staging[*blockHash] = clone
-	return nil
-}
-
-func (bms *blockStore) IsStaged() bool {
-	return len(bms.staging) != 0 || len(bms.toDelete) != 0
-}
-
-func (bms *blockStore) Discard() {
-	bms.staging = make(map[externalapi.DomainHash]*externalapi.DomainBlock)
-	bms.toDelete = make(map[externalapi.DomainHash]struct{})
-}
-
-func (bms *blockStore) Commit(dbTx model.DBTransaction) error {
-	for hash, block := range bms.staging {
-		blockBytes, err := bms.serializeBlock(block)
-		if err != nil {
-			return err
-		}
-		err = dbTx.Put(bms.hashAsKey(&hash), blockBytes)
-		if err != nil {
-			return err
-		}
-	}
-
-	for hash := range bms.toDelete {
-		err := dbTx.Delete(bms.hashAsKey(&hash))
-		if err != nil {
-			return err
-		}
-	}
-
-	bms.Discard()
-	return nil
-}
-
-// Block gets the block associated with the given blockHash
-func (bms *blockStore) Block(dbContext model.DBReader, blockHash *externalapi.DomainHash) (*externalapi.DomainBlock, error) {
-	if block, ok := bms.staging[*blockHash]; ok {
-		return block, nil
-	}
-
-	blockBytes, err := dbContext.Get(bms.hashAsKey(blockHash))
 	if err != nil {
 		return nil, err
 	}
 
-	return bms.deserializeBlock(blockBytes)
+	return blockStore, nil
+}
+
+func (bs *blockStore) initializeCount(dbContext model.DBReader) error {
+	count := uint64(0)
+	hasCountBytes, err := dbContext.Has(countKey)
+	if err != nil {
+		return err
+	}
+	if hasCountBytes {
+		countBytes, err := dbContext.Get(countKey)
+		if err != nil {
+			return err
+		}
+		count, err = bs.deserializeBlockCount(countBytes)
+		if err != nil {
+			return err
+		}
+	}
+	bs.count = count
+	return nil
+}
+
+// Stage stages the given block for the given blockHash
+func (bs *blockStore) Stage(blockHash *externalapi.DomainHash, block *externalapi.DomainBlock) error {
+	clone, err := bs.clone(block)
+	if err != nil {
+		return err
+	}
+
+	bs.staging[*blockHash] = clone
+	return nil
+}
+
+func (bs *blockStore) IsStaged() bool {
+	return len(bs.staging) != 0 || len(bs.toDelete) != 0
+}
+
+func (bs *blockStore) Discard() {
+	bs.staging = make(map[externalapi.DomainHash]*externalapi.DomainBlock)
+	bs.toDelete = make(map[externalapi.DomainHash]struct{})
+}
+
+func (bs *blockStore) Commit(dbTx model.DBTransaction) error {
+	for hash, block := range bs.staging {
+		blockBytes, err := bs.serializeBlock(block)
+		if err != nil {
+			return err
+		}
+		err = dbTx.Put(bs.hashAsKey(&hash), blockBytes)
+		if err != nil {
+			return err
+		}
+	}
+
+	for hash := range bs.toDelete {
+		err := dbTx.Delete(bs.hashAsKey(&hash))
+		if err != nil {
+			return err
+		}
+	}
+
+	err := bs.commitCount(dbTx)
+	if err != nil {
+		return err
+	}
+
+	bs.Discard()
+	return nil
+}
+
+// Block gets the block associated with the given blockHash
+func (bs *blockStore) Block(dbContext model.DBReader, blockHash *externalapi.DomainHash) (*externalapi.DomainBlock, error) {
+	if block, ok := bs.staging[*blockHash]; ok {
+		return block, nil
+	}
+
+	blockBytes, err := dbContext.Get(bs.hashAsKey(blockHash))
+	if err != nil {
+		return nil, err
+	}
+
+	return bs.deserializeBlock(blockBytes)
 }
 
 // HasBlock returns whether a block with a given hash exists in the store.
-func (bms *blockStore) HasBlock(dbContext model.DBReader, blockHash *externalapi.DomainHash) (bool, error) {
+func (bs *blockStore) HasBlock(dbContext model.DBReader, blockHash *externalapi.DomainHash) (bool, error) {
-	if _, ok := bms.staging[*blockHash]; ok {
+	if _, ok := bs.staging[*blockHash]; ok {
 		return true, nil
 	}
 
-	exists, err := dbContext.Has(bms.hashAsKey(blockHash))
+	exists, err := dbContext.Has(bs.hashAsKey(blockHash))
 	if err != nil {
 		return false, err
 	}
@@ -96,11 +130,11 @@ func (bms *blockStore) HasBlock(dbContext model.DBReader, blockHash *externalapi
 }
 
 // Blocks gets the blocks associated with the given blockHashes
-func (bms *blockStore) Blocks(dbContext model.DBReader, blockHashes []*externalapi.DomainHash) ([]*externalapi.DomainBlock, error) {
+func (bs *blockStore) Blocks(dbContext model.DBReader, blockHashes []*externalapi.DomainHash) ([]*externalapi.DomainBlock, error) {
 	blocks := make([]*externalapi.DomainBlock, len(blockHashes))
 	for i, hash := range blockHashes {
 		var err error
-		blocks[i], err = bms.Block(dbContext, hash)
+		blocks[i], err = bs.Block(dbContext, hash)
 		if err != nil {
 			return nil, err
 		}
@@ -109,20 +143,20 @@ func (bms *blockStore) Blocks(dbContext model.DBReader, blockHashes []*externala
 }
 
 // Delete deletes the block associated with the given blockHash
-func (bms *blockStore) Delete(blockHash *externalapi.DomainHash) {
+func (bs *blockStore) Delete(blockHash *externalapi.DomainHash) {
-	if _, ok := bms.staging[*blockHash]; ok {
+	if _, ok := bs.staging[*blockHash]; ok {
-		delete(bms.staging, *blockHash)
+		delete(bs.staging, *blockHash)
 		return
 	}
-	bms.toDelete[*blockHash] = struct{}{}
+	bs.toDelete[*blockHash] = struct{}{}
 }
 
-func (bms *blockStore) serializeBlock(block *externalapi.DomainBlock) ([]byte, error) {
+func (bs *blockStore) serializeBlock(block *externalapi.DomainBlock) ([]byte, error) {
 	dbBlock := serialization.DomainBlockToDbBlock(block)
 	return proto.Marshal(dbBlock)
 }
 
-func (bms *blockStore) deserializeBlock(blockBytes []byte) (*externalapi.DomainBlock, error) {
+func (bs *blockStore) deserializeBlock(blockBytes []byte) (*externalapi.DomainBlock, error) {
 	dbBlock := &serialization.DbBlock{}
 	err := proto.Unmarshal(blockBytes, dbBlock)
 	if err != nil {
@@ -131,27 +165,47 @@ func (bms *blockStore) deserializeBlock(blockBytes []byte) (*externalapi.DomainB
 	return serialization.DbBlockToDomainBlock(dbBlock)
 }
 
-func (bms *blockStore) hashAsKey(hash *externalapi.DomainHash) model.DBKey {
+func (bs *blockStore) hashAsKey(hash *externalapi.DomainHash) model.DBKey {
 	return bucket.Key(hash[:])
 }
 
-func (bms *blockStore) clone(block *externalapi.DomainBlock) (*externalapi.DomainBlock, error) {
+func (bs *blockStore) clone(block *externalapi.DomainBlock) (*externalapi.DomainBlock, error) {
-	serialized, err := bms.serializeBlock(block)
+	serialized, err := bs.serializeBlock(block)
 	if err != nil {
 		return nil, err
 	}
 
-	return bms.deserializeBlock(serialized)
+	return bs.deserializeBlock(serialized)
 }
 
-func (bms *blockStore) Count(dbContext model.DBReader) (uint64, error) {
+func (bs *blockStore) Count() uint64 {
-	cursor, err := dbContext.Cursor(bucket)
+	return bs.count + uint64(len(bs.staging)) - uint64(len(bs.toDelete))
+}
+
+func (bs *blockStore) deserializeBlockCount(countBytes []byte) (uint64, error) {
+	dbBlockCount := &serialization.DbBlockCount{}
+	err := proto.Unmarshal(countBytes, dbBlockCount)
 	if err != nil {
 		return 0, err
 	}
-	count := uint64(0)
+	return dbBlockCount.Count, nil
-	for cursor.Next() {
+}
-		count++
+
-	}
+func (bs *blockStore) commitCount(dbTx model.DBTransaction) error {
-	return count, nil
+	count := bs.Count()
+	countBytes, err := bs.serializeBlockCount(count)
+	if err != nil {
+		return err
+	}
+	err = dbTx.Put(countKey, countBytes)
+	if err != nil {
+		return err
+	}
+	bs.count = count
+	return nil
+}
+
+func (bs *blockStore) serializeBlockCount(count uint64) ([]byte, error) {
+	dbBlockCount := &serialization.DbBlockCount{Count: count}
+	return proto.Marshal(dbBlockCount)
 }

domain/consensus/factory.go

@@ -59,10 +59,18 @@ func NewFactory() Factory {
 
 // NewConsensus instantiates a new Consensus
 func (f *factory) NewConsensus(dagParams *dagconfig.Params, db infrastructuredatabase.Database) (externalapi.Consensus, error) {
+	dbManager := consensusdatabase.New(db)
+
 	// Data Structures
 	acceptanceDataStore := acceptancedatastore.New()
-	blockStore := blockstore.New()
+	blockStore, err := blockstore.New(dbManager)
-	blockHeaderStore := blockheaderstore.New()
+	if err != nil {
+		return nil, err
+	}
+	blockHeaderStore, err := blockheaderstore.New(dbManager)
+	if err != nil {
+		return nil, err
+	}
 	blockRelationStore := blockrelationstore.New()
 	blockStatusStore := blockstatusstore.New()
 	multisetStore := multisetstore.New()
@@ -73,8 +81,6 @@ func (f *factory) NewConsensus(dagParams *dagconfig.Params, db infrastructuredat
 	ghostdagDataStore := ghostdagdatastore.New()
 	headerTipsStore := headertipsstore.New()
 
-	dbManager := consensusdatabase.New(db)
-
 	// Processes
 	reachabilityManager := reachabilitymanager.New(
 		dbManager,

domain/consensus/model/interface_datastructures_blockheaderstore.go

@@ -11,5 +11,5 @@ type BlockHeaderStore interface {
 	HasBlockHeader(dbContext DBReader, blockHash *externalapi.DomainHash) (bool, error)
 	BlockHeaders(dbContext DBReader, blockHashes []*externalapi.DomainHash) ([]*externalapi.DomainBlockHeader, error)
 	Delete(blockHash *externalapi.DomainHash)
-	Count(dbContext DBReader) (uint64, error)
+	Count() uint64
 }

domain/consensus/model/interface_datastructures_blockstore.go

@@ -11,5 +11,5 @@ type BlockStore interface {
 	HasBlock(dbContext DBReader, blockHash *externalapi.DomainHash) (bool, error)
 	Blocks(dbContext DBReader, blockHashes []*externalapi.DomainHash) ([]*externalapi.DomainBlock, error)
 	Delete(blockHash *externalapi.DomainHash)
-	Count(dbContext DBReader) (uint64, error)
+	Count() uint64
 }

domain/consensus/processes/syncmanager/syncinfo.go

@@ -25,14 +25,8 @@ func (sm *syncManager) syncInfo() (*externalapi.SyncInfo, error) {
 		}
 	}
 
-	headerCount, err := sm.getHeaderCount()
+	headerCount := sm.getHeaderCount()
-	if err != nil {
+	blockCount := sm.getBlockCount()
-		return nil, err
-	}
-	blockCount, err := sm.getBlockCount()
-	if err != nil {
-		return nil, err
-	}
 
 	return &externalapi.SyncInfo{
 		State:                syncState,
@@ -123,10 +117,10 @@ func (sm *syncManager) areHeaderTipsSynced(headerVirtualSelectedParentHash *exte
 	return timeDifference <= maxTimeDifference, nil
 }
 
-func (sm *syncManager) getHeaderCount() (uint64, error) {
+func (sm *syncManager) getHeaderCount() uint64 {
-	return sm.blockHeaderStore.Count(sm.databaseContext)
+	return sm.blockHeaderStore.Count()
 }
 
-func (sm *syncManager) getBlockCount() (uint64, error) {
+func (sm *syncManager) getBlockCount() uint64 {
-	return sm.blockStore.Count(sm.databaseContext)
+	return sm.blockStore.Count()
 }