forked from flyteorg/flyte
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[flyteadmin] Refactor panic recovery into middleware (flyteorg#5546)
* Refactor panic handling to middleware Signed-off-by: Jason Parraga <[email protected]> * Remove registration of old panicCounter Signed-off-by: Jason Parraga <[email protected]> * Add test coverage Signed-off-by: Jason Parraga <[email protected]> --------- Signed-off-by: Jason Parraga <[email protected]> Signed-off-by: Bugra Gedik <[email protected]>
- Loading branch information
1 parent
7306683
commit dd2957b
Showing
17 changed files
with
177 additions
and
119 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
61 changes: 61 additions & 0 deletions
61
flyteadmin/pkg/rpc/adminservice/middleware/recovery_interceptor.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,61 @@ | ||
package middleware | ||
|
||
import ( | ||
"context" | ||
"runtime/debug" | ||
|
||
"github.com/prometheus/client_golang/prometheus" | ||
"google.golang.org/grpc" | ||
"google.golang.org/grpc/codes" | ||
"google.golang.org/grpc/status" | ||
|
||
"github.com/flyteorg/flyte/flytestdlib/logger" | ||
"github.com/flyteorg/flyte/flytestdlib/promutils" | ||
) | ||
|
||
// RecoveryInterceptor is a struct for creating gRPC interceptors that handle panics in go | ||
type RecoveryInterceptor struct { | ||
panicCounter prometheus.Counter | ||
} | ||
|
||
// NewRecoveryInterceptor creates a new RecoveryInterceptor with metrics under the provided scope | ||
func NewRecoveryInterceptor(adminScope promutils.Scope) *RecoveryInterceptor { | ||
panicCounter := adminScope.MustNewCounter("handler_panic", "panics encountered while handling gRPC requests") | ||
return &RecoveryInterceptor{ | ||
panicCounter: panicCounter, | ||
} | ||
} | ||
|
||
// UnaryServerInterceptor returns a new unary server interceptor for panic recovery. | ||
func (ri *RecoveryInterceptor) UnaryServerInterceptor() grpc.UnaryServerInterceptor { | ||
return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (_ any, err error) { | ||
|
||
defer func() { | ||
if r := recover(); r != nil { | ||
ri.panicCounter.Inc() | ||
logger.Errorf(ctx, "panic-ed for request: [%+v] to %s with err: %v with Stack: %v", req, info.FullMethod, r, string(debug.Stack())) | ||
// Return INTERNAL to client with no info as to not leak implementation details | ||
err = status.Errorf(codes.Internal, "") | ||
} | ||
}() | ||
|
||
return handler(ctx, req) | ||
} | ||
} | ||
|
||
// StreamServerInterceptor returns a new streaming server interceptor for panic recovery. | ||
func (ri *RecoveryInterceptor) StreamServerInterceptor() grpc.StreamServerInterceptor { | ||
return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) { | ||
|
||
defer func() { | ||
if r := recover(); r != nil { | ||
ri.panicCounter.Inc() | ||
logger.Errorf(stream.Context(), "panic-ed for stream to %s with err: %v with Stack: %v", info.FullMethod, r, string(debug.Stack())) | ||
// Return INTERNAL to client with no info as to not leak implementation details | ||
err = status.Errorf(codes.Internal, "") | ||
} | ||
}() | ||
|
||
return handler(srv, stream) | ||
} | ||
} |
Oops, something went wrong.