Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
orca: allow a ServerMetricsProvider to be passed to the ORCA service …
…and ServerOption (#6223)
- Loading branch information
Showing
9 changed files
with
546 additions
and
344 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 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,196 @@ | ||
| /* | ||
| * | ||
| * Copyright 2022 gRPC authors. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| * | ||
| */ | ||
|
|
||
| package orca | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync" | ||
|
|
||
| "google.golang.org/grpc" | ||
| grpcinternal "google.golang.org/grpc/internal" | ||
| "google.golang.org/grpc/metadata" | ||
| "google.golang.org/grpc/orca/internal" | ||
| "google.golang.org/protobuf/proto" | ||
| ) | ||
|
|
||
| // CallMetricsRecorder allows a service method handler to record per-RPC | ||
| // metrics. It contains all utilization-based metrics from | ||
| // ServerMetricsRecorder as well as additional request cost metrics. | ||
| type CallMetricsRecorder interface { | ||
| ServerMetricsRecorder | ||
|
|
||
| // SetRequestCost sets the relevant server metric. | ||
| SetRequestCost(name string, val float64) | ||
| // DeleteRequestCost deletes the relevant server metric to prevent it | ||
| // from being sent. | ||
| DeleteRequestCost(name string) | ||
|
|
||
| // SetNamedMetric sets the relevant server metric. | ||
| SetNamedMetric(name string, val float64) | ||
| // DeleteNamedMetric deletes the relevant server metric to prevent it | ||
| // from being sent. | ||
| DeleteNamedMetric(name string) | ||
| } | ||
|
|
||
| type callMetricsRecorderCtxKey struct{} | ||
|
|
||
| // CallMetricsRecorderFromContext returns the RPC-specific custom metrics | ||
| // recorder embedded in the provided RPC context. | ||
| // | ||
| // Returns nil if no custom metrics recorder is found in the provided context, | ||
| // which will be the case when custom metrics reporting is not enabled. | ||
| func CallMetricsRecorderFromContext(ctx context.Context) CallMetricsRecorder { | ||
| rw, ok := ctx.Value(callMetricsRecorderCtxKey{}).(*recorderWrapper) | ||
| if !ok { | ||
| return nil | ||
| } | ||
| return rw.recorder() | ||
| } | ||
|
|
||
| // recorderWrapper is a wrapper around a CallMetricsRecorder to ensure that | ||
| // concurrent calls to CallMetricsRecorderFromContext() results in only one | ||
| // allocation of the underlying metrics recorder, while also allowing for lazy | ||
| // initialization of the recorder itself. | ||
| type recorderWrapper struct { | ||
| once sync.Once | ||
| r CallMetricsRecorder | ||
| smp ServerMetricsProvider | ||
| } | ||
|
|
||
| func (rw *recorderWrapper) recorder() CallMetricsRecorder { | ||
| rw.once.Do(func() { | ||
| rw.r = newServerMetricsRecorder() | ||
| }) | ||
| return rw.r | ||
| } | ||
|
|
||
| // setTrailerMetadata adds a trailer metadata entry with key being set to | ||
| // `internal.TrailerMetadataKey` and value being set to the binary-encoded | ||
| // orca.OrcaLoadReport protobuf message. | ||
| // | ||
| // This function is called from the unary and streaming interceptors defined | ||
| // above. Any errors encountered here are not propagated to the caller because | ||
| // they are ignored there. Hence we simply log any errors encountered here at | ||
| // warning level, and return nothing. | ||
| func (rw *recorderWrapper) setTrailerMetadata(ctx context.Context) { | ||
| var sm *ServerMetrics | ||
| if rw.smp != nil { | ||
| sm = rw.smp.ServerMetrics() | ||
| sm.merge(rw.r.ServerMetrics()) | ||
| } else { | ||
| sm = rw.r.ServerMetrics() | ||
| } | ||
|
|
||
| b, err := proto.Marshal(sm.toLoadReportProto()) | ||
| if err != nil { | ||
| logger.Warningf("Failed to marshal load report: %v", err) | ||
| return | ||
| } | ||
| if err := grpc.SetTrailer(ctx, metadata.Pairs(internal.TrailerMetadataKey, string(b))); err != nil { | ||
| logger.Warningf("Failed to set trailer metadata: %v", err) | ||
| } | ||
| } | ||
|
|
||
| var joinServerOptions = grpcinternal.JoinServerOptions.(func(...grpc.ServerOption) grpc.ServerOption) | ||
|
|
||
| // CallMetricsServerOption returns a server option which enables the reporting | ||
| // of per-RPC custom backend metrics for unary and streaming RPCs. | ||
| // | ||
| // Server applications interested in injecting custom backend metrics should | ||
| // pass the server option returned from this function as the first argument to | ||
| // grpc.NewServer(). | ||
| // | ||
| // Subsequently, server RPC handlers can retrieve a reference to the RPC | ||
| // specific custom metrics recorder [CallMetricsRecorder] to be used, via a call | ||
| // to CallMetricsRecorderFromContext(), and inject custom metrics at any time | ||
| // during the RPC lifecycle. | ||
| // | ||
| // The injected custom metrics will be sent as part of trailer metadata, as a | ||
| // binary-encoded [ORCA LoadReport] protobuf message, with the metadata key | ||
| // being set be "endpoint-load-metrics-bin". | ||
| // | ||
| // If a non-nil ServerMetricsProvider is provided, the gRPC server will | ||
| // transmit the metrics it provides, overwritten by any per-RPC metrics given | ||
| // to the CallMetricsRecorder. A ServerMetricsProvider is typically obtained | ||
| // by calling NewServerMetricsRecorder. | ||
| // | ||
| // [ORCA LoadReport]: https://github.com/cncf/xds/blob/main/xds/data/orca/v3/orca_load_report.proto#L15 | ||
| func CallMetricsServerOption(smp ServerMetricsProvider) grpc.ServerOption { | ||
| return joinServerOptions(grpc.ChainUnaryInterceptor(unaryInt(smp)), grpc.ChainStreamInterceptor(streamInt(smp))) | ||
| } | ||
|
|
||
| func unaryInt(smp ServerMetricsProvider) func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { | ||
| return func(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { | ||
| // We don't allocate the metric recorder here. It will be allocated the | ||
| // first time the user calls CallMetricsRecorderFromContext(). | ||
| rw := &recorderWrapper{smp: smp} | ||
| ctxWithRecorder := newContextWithRecorderWrapper(ctx, rw) | ||
|
|
||
| resp, err := handler(ctxWithRecorder, req) | ||
|
|
||
| // It is safe to access the underlying metric recorder inside the wrapper at | ||
| // this point, as the user's RPC handler is done executing, and therefore | ||
| // there will be no more calls to CallMetricsRecorderFromContext(), which is | ||
| // where the metric recorder is lazy allocated. | ||
| if rw.r != nil { | ||
| rw.setTrailerMetadata(ctx) | ||
| } | ||
| return resp, err | ||
| } | ||
| } | ||
|
|
||
| func streamInt(smp ServerMetricsProvider) func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { | ||
| return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { | ||
| // We don't allocate the metric recorder here. It will be allocated the | ||
| // first time the user calls CallMetricsRecorderFromContext(). | ||
| rw := &recorderWrapper{smp: smp} | ||
| ws := &wrappedStream{ | ||
| ServerStream: ss, | ||
| ctx: newContextWithRecorderWrapper(ss.Context(), rw), | ||
| } | ||
|
|
||
| err := handler(srv, ws) | ||
|
|
||
| // It is safe to access the underlying metric recorder inside the wrapper at | ||
| // this point, as the user's RPC handler is done executing, and therefore | ||
| // there will be no more calls to CallMetricsRecorderFromContext(), which is | ||
| // where the metric recorder is lazy allocated. | ||
| if rw.r != nil { | ||
| rw.setTrailerMetadata(ss.Context()) | ||
| } | ||
| return err | ||
| } | ||
| } | ||
|
|
||
| func newContextWithRecorderWrapper(ctx context.Context, r *recorderWrapper) context.Context { | ||
| return context.WithValue(ctx, callMetricsRecorderCtxKey{}, r) | ||
| } | ||
|
|
||
| // wrappedStream wraps the grpc.ServerStream received by the streaming | ||
| // interceptor. Overrides only the Context() method to return a context which | ||
| // contains a reference to the CallMetricsRecorder corresponding to this | ||
| // stream. | ||
| type wrappedStream struct { | ||
| grpc.ServerStream | ||
| ctx context.Context | ||
| } | ||
|
|
||
| func (w *wrappedStream) Context() context.Context { | ||
| return w.ctx | ||
| } |
Oops, something went wrong.