Update waiting validator methods to be context aware (#8078)

* Update waiting validator methods to be context aware

* Enable debug logging for validator in e2e

* invert logic in deferred function, remove for loop

* return early if wait is 0 or less

* overwrite ctx

* t.Stop() chan closure is hanging. Doesnt make sense, so ignoring that cleanup task. It shouldnt happen at runtime anyway

Co-authored-by: terence tsao <terence@prysmaticlabs.com>

endtoend/components/validator.go

@@ -77,6 +77,7 @@ func StartNewValidatorClient(t *testing.T, config *types.E2EConfig, validatorNum
 		"--force-clear-db",
 		"--e2e-config",
 		"--accept-terms-of-use",
+		"--verbosity=debug",
 	}
 	args = append(args, featureconfig.E2EValidatorFlags...)
 	args = append(args, config.ValidatorFlags...)

validator/client/BUILD.bazel

@@ -34,6 +34,7 @@ go_library(
         "//shared/rand:go_default_library",
         "//shared/slotutil:go_default_library",
         "//shared/timeutils:go_default_library",
+        "//shared/traceutil:go_default_library",
         "//validator/accounts/wallet:go_default_library",
         "//validator/db:go_default_library",
         "//validator/db/kv:go_default_library",

validator/client/aggregate.go

@@ -12,6 +12,7 @@ import (
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/slotutil"
 	"github.com/prysmaticlabs/prysm/shared/timeutils"
+	"github.com/prysmaticlabs/prysm/shared/traceutil"
 	"go.opencensus.io/trace"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -143,7 +144,7 @@ func (v *validator) signSlot(ctx context.Context, pubKey [48]byte, slot uint64)
 // such that any attestations from this slot have time to reach the beacon node
 // before creating the aggregated attestation.
 func (v *validator) waitToSlotTwoThirds(ctx context.Context, slot uint64) {
-	_, span := trace.StartSpan(ctx, "validator.waitToSlotTwoThirds")
+	ctx, span := trace.StartSpan(ctx, "validator.waitToSlotTwoThirds")
 	defer span.End()
 
 	oneThird := slotutil.DivideSlotBy(3 /* one third of slot duration */)
@@ -152,7 +153,19 @@ func (v *validator) waitToSlotTwoThirds(ctx context.Context, slot uint64) {
 
 	startTime := slotutil.SlotStartTime(v.genesisTime, slot)
 	finalTime := startTime.Add(delay)
-	time.Sleep(timeutils.Until(finalTime))
+	wait := timeutils.Until(finalTime)
+	if wait <= 0 {
+		return
+	}
+	t := time.NewTimer(wait)
+	defer t.Stop()
+	select {
+	case <-ctx.Done():
+		traceutil.AnnotateError(span, ctx.Err())
+		return
+	case <-t.C:
+		return
+	}
 }
 
 // This returns the signature of validator signing over aggregate and

validator/client/aggregate_test.go

@@ -140,6 +140,21 @@ func TestWaitForSlotTwoThird_WaitCorrectly(t *testing.T) {
 	assert.Equal(t, twoThirdTime.Unix(), currentTime.Unix())
 }
 
+func TestWaitForSlotTwoThird_DoneContext_ReturnsImmediately(t *testing.T) {
+	validator, _, _, finish := setup(t)
+	defer finish()
+	currentTime := timeutils.Now()
+	numOfSlots := uint64(4)
+	validator.genesisTime = uint64(currentTime.Unix()) - (numOfSlots * params.BeaconConfig().SecondsPerSlot)
+
+	expectedTime := timeutils.Now()
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	validator.waitToSlotTwoThirds(ctx, numOfSlots)
+	currentTime = timeutils.Now()
+	assert.Equal(t, expectedTime.Unix(), currentTime.Unix())
+}
+
 func TestAggregateAndProofSignature_CanSignValidSignature(t *testing.T) {
 	validator, m, validatorKey, finish := setup(t)
 	defer finish()

validator/client/attest.go

@@ -16,6 +16,7 @@ import (
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/slotutil"
 	"github.com/prysmaticlabs/prysm/shared/timeutils"
+	"github.com/prysmaticlabs/prysm/shared/traceutil"
 	"github.com/sirupsen/logrus"
 	"go.opencensus.io/trace"
 )
@@ -218,13 +219,25 @@ func (v *validator) saveAttesterIndexToData(data *ethpb.AttestationData, index u
 // waitToSlotOneThird waits until one third through the current slot period
 // such that head block for beacon node can get updated.
 func (v *validator) waitToSlotOneThird(ctx context.Context, slot uint64) {
-	_, span := trace.StartSpan(ctx, "validator.waitToSlotOneThird")
+	ctx, span := trace.StartSpan(ctx, "validator.waitToSlotOneThird")
 	defer span.End()
 
 	delay := slotutil.DivideSlotBy(3 /* a third of the slot duration */)
 	startTime := slotutil.SlotStartTime(v.genesisTime, slot)
 	finalTime := startTime.Add(delay)
-	time.Sleep(timeutils.Until(finalTime))
+	wait := timeutils.Until(finalTime)
+	if wait <= 0 {
+		return
+	}
+	t := time.NewTimer(wait)
+	defer t.Stop()
+	select {
+	case <-ctx.Done():
+		traceutil.AnnotateError(span, ctx.Err())
+		return
+	case <-t.C:
+		return
+	}
 }
 
 func attestationLogFields(pubKey [48]byte, indexedAtt *ethpb.IndexedAttestation) logrus.Fields {