Add span to HTR and skip slot cache (#5197)

* Add a bit more span data
* missing import
* Merge branch 'master' into more-spans

beacon-chain/blockchain/service.go

@@ -314,7 +314,7 @@ func (s *Service) saveGenesisValidators(ctx context.Context, state *stateTrie.Be
 
 // This gets called when beacon chain is first initialized to save genesis data (state, block, and more) in db.
 func (s *Service) saveGenesisData(ctx context.Context, genesisState *stateTrie.BeaconState) error {
-	stateRoot, err := genesisState.HashTreeRoot()
+	stateRoot, err := genesisState.HashTreeRoot(ctx)
 	if err != nil {
 		return err
 	}

beacon-chain/cache/BUILD.bazel

@@ -25,6 +25,7 @@ go_library(
         "@com_github_prometheus_client_golang//prometheus/promauto:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
         "@io_k8s_client_go//tools/cache:go_default_library",
+        "@io_opencensus_go//trace:go_default_library",
     ],
 )
 

beacon-chain/cache/skip_slot_cache.go

@@ -11,6 +11,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	stateTrie "github.com/prysmaticlabs/prysm/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/shared/featureconfig"
+	"go.opencensus.io/trace"
 )
 
 var (
@@ -47,6 +48,8 @@ func NewSkipSlotCache() *SkipSlotCache {
 // Get waits for any in progress calculation to complete before returning a
 // cached response, if any.
 func (c *SkipSlotCache) Get(ctx context.Context, slot uint64) (*stateTrie.BeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "skipSlotCache.Get")
+	defer span.End()
 	if !featureconfig.Get().EnableSkipSlotsCache {
 		// Return a miss result if cache is not enabled.
 		skipSlotCacheMiss.Inc()
@@ -57,6 +60,7 @@ func (c *SkipSlotCache) Get(ctx context.Context, slot uint64) (*stateTrie.Beacon
 
 	// Another identical request may be in progress already. Let's wait until
 	// any in progress request resolves or our timeout is exceeded.
+	inProgress := false
 	for {
 		if ctx.Err() != nil {
 			return nil, ctx.Err()
@@ -67,6 +71,7 @@ func (c *SkipSlotCache) Get(ctx context.Context, slot uint64) (*stateTrie.Beacon
 			c.lock.RUnlock()
 			break
 		}
+		inProgress = true
 		c.lock.RUnlock()
 
 		// This increasing backoff is to decrease the CPU cycles while waiting
@@ -75,14 +80,17 @@ func (c *SkipSlotCache) Get(ctx context.Context, slot uint64) (*stateTrie.Beacon
 		delay *= delayFactor
 		delay = math.Min(delay, maxDelay)
 	}
+	span.AddAttributes(trace.BoolAttribute("inProgress", inProgress))
 
 	item, exists := c.cache.Get(slot)
 
 	if exists && item != nil {
 		skipSlotCacheHit.Inc()
+		span.AddAttributes(trace.BoolAttribute("hit", true))
 		return item.(*stateTrie.BeaconState).Copy(), nil
 	}
 	skipSlotCacheMiss.Inc()
+	span.AddAttributes(trace.BoolAttribute("hit", false))
 	return nil, nil
 }
 

beacon-chain/core/state/benchmarks_test.go

@@ -136,14 +136,14 @@ func BenchmarkHashTreeRootState_FullState(b *testing.B) {
 	}
 
 	// Hydrate the HashTreeRootState cache.
-	if _, err := beaconState.HashTreeRoot(); err != nil {
+	if _, err := beaconState.HashTreeRoot(ctx); err != nil {
 		b.Fatal(err)
 	}
 
 	b.N = 50
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		if _, err := beaconState.HashTreeRoot(); err != nil {
+		if _, err := beaconState.HashTreeRoot(ctx); err != nil {
 			b.Fatal(err)
 		}
 	}

beacon-chain/core/state/transition.go

@@ -72,7 +72,7 @@ func ExecuteStateTransition(
 	interop.WriteBlockToDisk(signed, false)
 	interop.WriteStateToDisk(state)
 
-	postStateRoot, err := state.HashTreeRoot()
+	postStateRoot, err := state.HashTreeRoot(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -181,7 +181,7 @@ func CalculateStateRoot(
 		return [32]byte{}, errors.Wrap(err, "could not process block")
 	}
 
-	return state.HashTreeRoot()
+	return state.HashTreeRoot(ctx)
 }
 
 // ProcessSlot happens every slot and focuses on the slot counter and block roots record updates.
@@ -205,7 +205,7 @@ func ProcessSlot(ctx context.Context, state *stateTrie.BeaconState) (*stateTrie.
 	defer span.End()
 	span.AddAttributes(trace.Int64Attribute("slot", int64(state.Slot())))
 
-	prevStateRoot, err := state.HashTreeRoot()
+	prevStateRoot, err := state.HashTreeRoot(ctx)
 	if err != nil {
 		return nil, err
 	}

beacon-chain/interop-cold-start/service.go

@@ -153,7 +153,7 @@ func (s *Service) DepositsNumberAndRootAtHeight(ctx context.Context, blockHeight
 
 func (s *Service) saveGenesisState(ctx context.Context, genesisState *stateTrie.BeaconState) error {
 	s.chainStartDeposits = make([]*ethpb.Deposit, genesisState.NumValidators())
-	stateRoot, err := genesisState.HashTreeRoot()
+	stateRoot, err := genesisState.HashTreeRoot(ctx)
 	if err != nil {
 		return err
 	}

beacon-chain/rpc/validator/proposer_test.go

@@ -46,7 +46,7 @@ func TestGetBlock_OK(t *testing.T) {
 
 	beaconState, privKeys := testutil.DeterministicGenesisState(t, params.BeaconConfig().MinGenesisActiveValidatorCount)
 
-	stateRoot, err := beaconState.HashTreeRoot()
+	stateRoot, err := beaconState.HashTreeRoot(ctx)
 	if err != nil {
 		t.Fatalf("Could not hash genesis state: %v", err)
 	}
@@ -151,7 +151,7 @@ func TestGetBlock_AddsUnaggregatedAtts(t *testing.T) {
 
 	beaconState, privKeys := testutil.DeterministicGenesisState(t, params.BeaconConfig().MinGenesisActiveValidatorCount)
 
-	stateRoot, err := beaconState.HashTreeRoot()
+	stateRoot, err := beaconState.HashTreeRoot(ctx)
 	if err != nil {
 		t.Fatalf("Could not hash genesis state: %v", err)
 	}
@@ -319,7 +319,7 @@ func TestComputeStateRoot_OK(t *testing.T) {
 
 	beaconState, privKeys := testutil.DeterministicGenesisState(t, 100)
 
-	stateRoot, err := beaconState.HashTreeRoot()
+	stateRoot, err := beaconState.HashTreeRoot(ctx)
 	if err != nil {
 		t.Fatalf("Could not hash genesis state: %v", err)
 	}

beacon-chain/state/BUILD.bazel

@@ -32,6 +32,7 @@ go_library(
         "@com_github_protolambda_zssz//merkle:go_default_library",
         "@com_github_prysmaticlabs_ethereumapis//eth/v1alpha1:go_default_library",
         "@com_github_prysmaticlabs_go_bitfield//:go_default_library",
+        "@io_opencensus_go//trace:go_default_library",
     ],
 )
 

beacon-chain/state/state_trie.go

@@ -1,12 +1,11 @@
 package state
 
 import (
+	"context"
 	"runtime"
 	"sort"
 	"sync"
 
-	"github.com/prysmaticlabs/prysm/shared/sliceutil"
-
 	"github.com/gogo/protobuf/proto"
 	"github.com/pkg/errors"
 	"github.com/protolambda/zssz/merkle"
@@ -18,6 +17,8 @@ import (
 	"github.com/prysmaticlabs/prysm/shared/hashutil"
 	"github.com/prysmaticlabs/prysm/shared/memorypool"
 	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/prysmaticlabs/prysm/shared/sliceutil"
+	"go.opencensus.io/trace"
 )
 
 // InitializeFromProto the beacon state from a protobuf representation.
@@ -178,7 +179,10 @@ func (b *BeaconState) Copy() *BeaconState {
 
 // HashTreeRoot of the beacon state retrieves the Merkle root of the trie
 // representation of the beacon state based on the eth2 Simple Serialize specification.
-func (b *BeaconState) HashTreeRoot() ([32]byte, error) {
+func (b *BeaconState) HashTreeRoot(ctx context.Context) ([32]byte, error) {
+	_, span := trace.StartSpan(ctx, "beaconState.HashTreeRoot")
+	defer span.End()
+
 	b.lock.Lock()
 	defer b.lock.Unlock()
 

beacon-chain/state/types_test.go

@@ -1,6 +1,7 @@
 package state_test
 
 import (
+	"context"
 	"reflect"
 	"strconv"
 	"testing"
@@ -18,6 +19,7 @@ import (
 
 func TestBeaconState_ProtoBeaconStateCompatibility(t *testing.T) {
 	params.UseMinimalConfig()
+	ctx := context.Background()
 	genesis := setupGenesisState(t, 64)
 	customState, err := stateTrie.InitializeFromProto(genesis)
 	if err != nil {
@@ -29,7 +31,7 @@ func TestBeaconState_ProtoBeaconStateCompatibility(t *testing.T) {
 		t.Fatal("Cloned states did not match")
 	}
 
-	r1, err := customState.HashTreeRoot()
+	r1, err := customState.HashTreeRoot(ctx)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -47,7 +49,7 @@ func TestBeaconState_ProtoBeaconStateCompatibility(t *testing.T) {
 	if err := customState.SetBalances(balances); err != nil {
 		t.Fatal(err)
 	}
-	r1, err = customState.HashTreeRoot()
+	r1, err = customState.HashTreeRoot(ctx)
 	if err != nil {
 		t.Fatal(err)
 	}

shared/testutil/block.go

@@ -49,6 +49,7 @@ func GenerateFullBlock(
 	conf *BlockGenConfig,
 	slot uint64,
 ) (*ethpb.SignedBeaconBlock, error) {
+	ctx := context.Background()
 	currentSlot := bState.Slot()
 	if currentSlot > slot {
 		return nil, fmt.Errorf("current slot in state is larger than given slot. %d > %d", currentSlot, slot)
@@ -106,7 +107,7 @@ func GenerateFullBlock(
 	}
 
 	newHeader := bState.LatestBlockHeader()
-	prevStateRoot, err := bState.HashTreeRoot()
+	prevStateRoot, err := bState.HashTreeRoot(ctx)
 	if err != nil {
 		return nil, err
 	}