chore(lib/common) implement byte pool to improve websocket subscription efficiency

dot/state/block.go

@@ -46,10 +46,12 @@ type BlockState struct {
 	lastFinalised common.Hash
 
 	// block notifiers
-	imported      map[byte]chan<- *types.Block
-	finalised     map[byte]chan<- *types.FinalisationInfo
-	importedLock  sync.RWMutex
-	finalisedLock sync.RWMutex
+	imported          map[byte]chan<- *types.Block
+	finalised         map[byte]chan<- *types.FinalisationInfo
+	importedLock      sync.RWMutex
+	finalisedLock     sync.RWMutex
+	importedBytepool  *common.BytePool
+	finalisedBytepool *common.BytePool
 
 	pruneKeyCh chan *types.Header
 }
@@ -80,6 +82,16 @@ func NewBlockState(db chaindb.Database, bt *blocktree.BlockTree) (*BlockState, e
 		return nil, fmt.Errorf("failed to get last finalised hash: %w", err)
 	}
 
+	bs.importedBytepool, err = common.NewBytePool256()
+	if err != nil {
+		logger.Debug("issue setting up imported byte pool", "error", err)
+	}
+
+	bs.finalisedBytepool, err = common.NewBytePool256()
+	if err != nil {
+		logger.Debug("issue setting up finalised byte pool", "error", err)
+	}
+
 	return bs, nil
 }
 
@@ -122,6 +134,16 @@ func NewBlockStateFromGenesis(db chaindb.Database, header *types.Header) (*Block
 		return nil, err
 	}
 
+	bs.importedBytepool, err = common.NewBytePool256()
+	if err != nil {
+		logger.Debug("issue setting up imported byte pool", "error", err)
+	}
+
+	bs.finalisedBytepool, err = common.NewBytePool256()
+	if err != nil {
+		logger.Debug("issue setting up finalised byte pool", "error", err)
+	}
+
 	return bs, nil
 }
 

dot/state/block_notify.go

@@ -17,9 +17,6 @@
 package state
 
 import (
-	"errors"
-	"math/rand"
-
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
 )
@@ -29,16 +26,9 @@ import (
 func (bs *BlockState) RegisterImportedChannel(ch chan<- *types.Block) (byte, error) {
 	bs.importedLock.RLock()
 
-	if len(bs.imported) == 256 {
-		return 0, errors.New("channel limit reached")
-	}
-
-	var id byte
-	for {
-		id = generateID()
-		if bs.imported[id] == nil {
-			break
-		}
+	id, err := bs.importedBytepool.Get()
+	if err != nil {
+		return 0, err
 	}
 
 	bs.importedLock.RUnlock()
@@ -54,16 +44,9 @@ func (bs *BlockState) RegisterImportedChannel(ch chan<- *types.Block) (byte, err
 func (bs *BlockState) RegisterFinalizedChannel(ch chan<- *types.FinalisationInfo) (byte, error) {
 	bs.finalisedLock.RLock()
 
-	if len(bs.finalised) == 256 {
-		return 0, errors.New("channel limit reached")
-	}
-
-	var id byte
-	for {
-		id = generateID()
-		if bs.finalised[id] == nil {
-			break
-		}
+	id, err := bs.finalisedBytepool.Get()
+	if err != nil {
+		return 0, err
 	}
 
 	bs.finalisedLock.RUnlock()
@@ -141,9 +124,3 @@ func (bs *BlockState) notifyFinalized(hash common.Hash, round, setID uint64) {
 		}(ch)
 	}
 }
-
-func generateID() byte {
-	// skipcq: GSC-G404
-	id := rand.Intn(256) //nolint
-	return byte(id)
-}

lib/common/bytepool.go

@@ -0,0 +1,66 @@
+// Copyright 2019 ChainSafe Systems (ON) Corp.
+// This file is part of gossamer.
+//
+// The gossamer library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The gossamer library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the gossamer library. If not, see <http://www.gnu.org/licenses/>.
+
+package common
+
+import "fmt"
+
+// BytePool struct to hold byte objects that will be contained in pool
+type BytePool struct {
+	c chan byte
+}
+
+// NewBytePool256 creates and initialises pool with 256 entries
+func NewBytePool256() (bp *BytePool, err error) {
+	bp = NewBytePool(256)
+	for i := 0; i < 256; i++ {
+		err = bp.Put(byte(i))
+	}
+	return
+}
+
+// NewBytePool creates a new empty byte pool with capacity of size
+func NewBytePool(size int) (bp *BytePool) {
+	return &BytePool{
+		c: make(chan byte, size),
+	}
+}
+
+// Get gets a Buffer from the BytePool, or creates a new one if none are
+// available in the pool.
+func (bp *BytePool) Get() (b byte, err error) {
+	select {
+	case b = <-bp.c:
+	default:
+		err = fmt.Errorf("all slots used")
+	}
+	return
+}
+
+// Put returns the given Buffer to the BytePool.
+func (bp *BytePool) Put(b byte) error {
+	select {
+	case bp.c <- b:
+		return nil
+	default: // Discard the buffer if the pool is full.
+		return fmt.Errorf("pool is full")
+	}
+}
+
+// NumPooled returns the number of items currently pooled.
+func (bp *BytePool) NumPooled() int {
+	return len(bp.c)
+}

lib/common/bytepool_test.go

@@ -0,0 +1,63 @@
+// Copyright 2019 ChainSafe Systems (ON) Corp.
+// This file is part of gossamer.
+//
+// The gossamer library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The gossamer library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the gossamer library. If not, see <http://www.gnu.org/licenses/>.
+
+package common
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestBytePool(t *testing.T) {
+	bp := NewBytePool(5)
+	require.Equal(t, 0, bp.NumPooled())
+
+	for i := 0; i < 5; i++ {
+		err := bp.Put(generateID())
+		require.NoError(t, err)
+	}
+	err := bp.Put(generateID())
+	require.EqualError(t, err, "pool is full")
+	require.Equal(t, 5, bp.NumPooled())
+
+	for i := 0; i < 5; i++ {
+		_, err := bp.Get() // nolint
+		require.NoError(t, err)
+	}
+	_, err = bp.Get()
+	require.EqualError(t, err, "all slots used")
+}
+
+func TestBytePool256(t *testing.T) {
+	bp, err := NewBytePool256()
+	require.NoError(t, err)
+	require.Equal(t, 256, bp.NumPooled())
+
+	for i := 0; i < 256; i++ {
+		_, err := bp.Get() // nolint
+		require.NoError(t, err)
+	}
+	_, err = bp.Get()
+	require.EqualError(t, err, "all slots used")
+}
+
+func generateID() byte {
+	// skipcq: GSC-G404
+	id := rand.Intn(256) //nolint
+	return byte(id)
+}