go/beacon/api: Support fetching next block epoch

.changelog/6409.feature.md

@@ -0,0 +1 @@
+go/beacon/api: Support fetching next block epoch

go/beacon/api/api.go

@@ -71,10 +71,17 @@ type Backend interface {
 	GetBaseEpoch(context.Context) (EpochTime, error)
 
 	// GetEpoch returns the epoch number at the specified block height.
+	//
 	// Calling this method with height `consensus.HeightLatest`, should
-	// return the epoch of latest known block.
+	// return the epoch of the latest finalized block.
 	GetEpoch(context.Context, int64) (EpochTime, error)
 
+	// GetNextEpoch returns the epoch number after the specified block height.
+	//
+	// Calling this method with height `consensus.HeightLatest`, should
+	// return the epoch of the next block after the latest finalized block.
+	GetNextEpoch(context.Context, int64) (EpochTime, error)
+
 	// GetFutureEpoch returns any future epoch that is currently scheduled
 	// to occur at a specific height.
 	//
@@ -108,6 +115,7 @@ type Backend interface {
 	WatchLatestEpoch(ctx context.Context) (<-chan EpochTime, pubsub.ClosableSubscription, error)
 
 	// GetBeacon gets the beacon for the provided block height.
+	//
 	// Calling this method with height `consensus.HeightLatest` should
 	// return the beacon for the latest finalized block.
 	GetBeacon(context.Context, int64) ([]byte, error)

go/beacon/api/grpc.go

@@ -18,6 +18,8 @@ var (
 	methodGetBaseEpoch = serviceName.NewMethod("GetBaseEpoch", nil)
 	// methodGetEpoch is the GetEpoch method.
 	methodGetEpoch = serviceName.NewMethod("GetEpoch", int64(0))
+	// methodGetNextEpoch is the GetNextEpoch method.
+	methodGetNextEpoch = serviceName.NewMethod("GetNextEpoch", int64(0))
 	// methodGetFutureEpoch is the GetFutureEpoch method.
 	methodGetFutureEpoch = serviceName.NewMethod("GetFutureEpoch", int64(0))
 	// methodGetEpochBlock is the GetEpochBlock method.
@@ -47,6 +49,10 @@ var (
 				MethodName: methodGetEpoch.ShortName(),
 				Handler:    handlerGetEpoch,
 			},
+			{
+				MethodName: methodGetNextEpoch.ShortName(),
+				Handler:    handlerGetNextEpoch,
+			},
 			{
 				MethodName: methodGetFutureEpoch.ShortName(),
 				Handler:    handlerGetFutureEpoch,
@@ -124,6 +130,29 @@ func handlerGetEpoch(
 	return interceptor(ctx, height, info, handler)
 }
 
+func handlerGetNextEpoch(
+	srv any,
+	ctx context.Context,
+	dec func(any) error,
+	interceptor grpc.UnaryServerInterceptor,
+) (any, error) {
+	var height int64
+	if err := dec(&height); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(Backend).GetNextEpoch(ctx, height)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: methodGetNextEpoch.FullName(),
+	}
+	handler := func(ctx context.Context, req any) (any, error) {
+		return srv.(Backend).GetNextEpoch(ctx, req.(int64))
+	}
+	return interceptor(ctx, height, info, handler)
+}
+
 func handlerGetFutureEpoch(
 	srv any,
 	ctx context.Context,
@@ -323,6 +352,14 @@ func (c *Client) GetEpoch(ctx context.Context, height int64) (EpochTime, error)
 	return rsp, nil
 }
 
+func (c *Client) GetNextEpoch(ctx context.Context, height int64) (EpochTime, error) {
+	var rsp EpochTime
+	if err := c.conn.Invoke(ctx, methodGetNextEpoch.FullName(), height, &rsp); err != nil {
+		return 0, err
+	}
+	return rsp, nil
+}
+
 func (c *Client) GetFutureEpoch(ctx context.Context, height int64) (*EpochTimeState, error) {
 	var rsp EpochTimeState
 	if err := c.conn.Invoke(ctx, methodGetFutureEpoch.FullName(), height, &rsp); err != nil {

go/consensus/cometbft/beacon/beacon.go

@@ -104,6 +104,33 @@ func (sc *ServiceClient) GetEpoch(ctx context.Context, height int64) (api.EpochT
 	return epoch, err
 }
 
+func (sc *ServiceClient) GetNextEpoch(ctx context.Context, height int64) (api.EpochTime, error) {
+	if height == consensus.HeightLatest {
+		latest, err := sc.consensus.GetLatestHeight(ctx)
+		if err != nil {
+			return api.EpochInvalid, err
+		}
+		height = latest
+	}
+
+	q, err := sc.querier.QueryAt(ctx, height)
+	if err != nil {
+		return api.EpochInvalid, err
+	}
+
+	future, err := q.FutureEpoch(ctx)
+	if err != nil {
+		return api.EpochInvalid, err
+	}
+
+	if future != nil && future.Height == height+1 {
+		return future.Epoch, nil
+	}
+
+	epoch, _, err := q.Epoch(ctx)
+	return epoch, err
+}
+
 func (sc *ServiceClient) GetFutureEpoch(ctx context.Context, height int64) (*api.EpochTimeState, error) {
 	q, err := sc.querier.QueryAt(ctx, height)
 	if err != nil {

go/worker/common/committee/node.go

@@ -244,9 +244,9 @@ func (n *Node) handleSuspend() {
 
 func (n *Node) updateHostedRuntimeVersion(rt *registry.Runtime) {
 	// Always take the latest epoch to avoid reverting to stale state.
-	epoch, err := n.Consensus.Beacon().GetEpoch(n.ctx, consensus.HeightLatest)
+	epoch, err := n.Consensus.Beacon().GetNextEpoch(n.ctx, consensus.HeightLatest)
 	if err != nil {
-		n.logger.Error("failed to fetch current epoch",
+		n.logger.Error("failed to fetch next block epoch",
 			"err", err,
 		)
 		return