mirror of https://github.com/grpc/grpc.git
Merge pull request #17577 from markdroth/lb_trailing_metadata
Allow LB policies to intercept recv_trailing_metadatapull/17753/head
commit
10fa278660
14 changed files with 418 additions and 0 deletions
@ -0,0 +1,240 @@ |
|||||||
|
/*
|
||||||
|
* |
||||||
|
* Copyright 2018 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. |
||||||
|
* |
||||||
|
*/ |
||||||
|
|
||||||
|
#include "test/core/util/test_lb_policies.h" |
||||||
|
|
||||||
|
#include <string> |
||||||
|
|
||||||
|
#include "src/core/ext/filters/client_channel/lb_policy.h" |
||||||
|
#include "src/core/ext/filters/client_channel/lb_policy_registry.h" |
||||||
|
#include "src/core/lib/channel/channel_args.h" |
||||||
|
#include "src/core/lib/channel/channelz.h" |
||||||
|
#include "src/core/lib/debug/trace.h" |
||||||
|
#include "src/core/lib/gprpp/memory.h" |
||||||
|
#include "src/core/lib/gprpp/orphanable.h" |
||||||
|
#include "src/core/lib/gprpp/ref_counted_ptr.h" |
||||||
|
#include "src/core/lib/iomgr/closure.h" |
||||||
|
#include "src/core/lib/iomgr/combiner.h" |
||||||
|
#include "src/core/lib/iomgr/error.h" |
||||||
|
#include "src/core/lib/iomgr/pollset_set.h" |
||||||
|
#include "src/core/lib/json/json.h" |
||||||
|
#include "src/core/lib/transport/connectivity_state.h" |
||||||
|
|
||||||
|
namespace grpc_core { |
||||||
|
|
||||||
|
TraceFlag grpc_trace_forwarding_lb(false, "forwarding_lb"); |
||||||
|
|
||||||
|
namespace { |
||||||
|
|
||||||
|
//
|
||||||
|
// ForwardingLoadBalancingPolicy
|
||||||
|
//
|
||||||
|
|
||||||
|
// A minimal forwarding class to avoid implementing a standalone test LB.
|
||||||
|
class ForwardingLoadBalancingPolicy : public LoadBalancingPolicy { |
||||||
|
public: |
||||||
|
ForwardingLoadBalancingPolicy(const Args& args, |
||||||
|
const std::string& delegate_policy_name) |
||||||
|
: LoadBalancingPolicy(args) { |
||||||
|
delegate_ = LoadBalancingPolicyRegistry::CreateLoadBalancingPolicy( |
||||||
|
delegate_policy_name.c_str(), args); |
||||||
|
grpc_pollset_set_add_pollset_set(delegate_->interested_parties(), |
||||||
|
interested_parties()); |
||||||
|
// Give re-resolution closure to delegate.
|
||||||
|
GRPC_CLOSURE_INIT(&on_delegate_request_reresolution_, |
||||||
|
OnDelegateRequestReresolutionLocked, this, |
||||||
|
grpc_combiner_scheduler(combiner())); |
||||||
|
Ref().release(); // held by callback.
|
||||||
|
delegate_->SetReresolutionClosureLocked(&on_delegate_request_reresolution_); |
||||||
|
} |
||||||
|
|
||||||
|
~ForwardingLoadBalancingPolicy() override = default; |
||||||
|
|
||||||
|
void UpdateLocked(const grpc_channel_args& args, |
||||||
|
grpc_json* lb_config) override { |
||||||
|
delegate_->UpdateLocked(args, lb_config); |
||||||
|
} |
||||||
|
|
||||||
|
bool PickLocked(PickState* pick, grpc_error** error) override { |
||||||
|
return delegate_->PickLocked(pick, error); |
||||||
|
} |
||||||
|
|
||||||
|
void CancelPickLocked(PickState* pick, grpc_error* error) override { |
||||||
|
delegate_->CancelPickLocked(pick, error); |
||||||
|
} |
||||||
|
|
||||||
|
void CancelMatchingPicksLocked(uint32_t initial_metadata_flags_mask, |
||||||
|
uint32_t initial_metadata_flags_eq, |
||||||
|
grpc_error* error) override { |
||||||
|
delegate_->CancelMatchingPicksLocked(initial_metadata_flags_mask, |
||||||
|
initial_metadata_flags_eq, error); |
||||||
|
} |
||||||
|
|
||||||
|
void NotifyOnStateChangeLocked(grpc_connectivity_state* state, |
||||||
|
grpc_closure* closure) override { |
||||||
|
delegate_->NotifyOnStateChangeLocked(state, closure); |
||||||
|
} |
||||||
|
|
||||||
|
grpc_connectivity_state CheckConnectivityLocked( |
||||||
|
grpc_error** connectivity_error) override { |
||||||
|
return delegate_->CheckConnectivityLocked(connectivity_error); |
||||||
|
} |
||||||
|
|
||||||
|
void HandOffPendingPicksLocked(LoadBalancingPolicy* new_policy) override { |
||||||
|
delegate_->HandOffPendingPicksLocked(new_policy); |
||||||
|
} |
||||||
|
|
||||||
|
void ExitIdleLocked() override { delegate_->ExitIdleLocked(); } |
||||||
|
|
||||||
|
void ResetBackoffLocked() override { delegate_->ResetBackoffLocked(); } |
||||||
|
|
||||||
|
void FillChildRefsForChannelz( |
||||||
|
channelz::ChildRefsList* child_subchannels, |
||||||
|
channelz::ChildRefsList* child_channels) override { |
||||||
|
delegate_->FillChildRefsForChannelz(child_subchannels, child_channels); |
||||||
|
} |
||||||
|
|
||||||
|
private: |
||||||
|
void ShutdownLocked() override { |
||||||
|
delegate_.reset(); |
||||||
|
TryReresolutionLocked(&grpc_trace_forwarding_lb, GRPC_ERROR_CANCELLED); |
||||||
|
} |
||||||
|
|
||||||
|
static void OnDelegateRequestReresolutionLocked(void* arg, |
||||||
|
grpc_error* error) { |
||||||
|
ForwardingLoadBalancingPolicy* self = |
||||||
|
static_cast<ForwardingLoadBalancingPolicy*>(arg); |
||||||
|
if (error != GRPC_ERROR_NONE || self->delegate_ == nullptr) { |
||||||
|
self->Unref(); |
||||||
|
return; |
||||||
|
} |
||||||
|
self->TryReresolutionLocked(&grpc_trace_forwarding_lb, GRPC_ERROR_NONE); |
||||||
|
self->delegate_->SetReresolutionClosureLocked( |
||||||
|
&self->on_delegate_request_reresolution_); |
||||||
|
} |
||||||
|
|
||||||
|
OrphanablePtr<LoadBalancingPolicy> delegate_; |
||||||
|
grpc_closure on_delegate_request_reresolution_; |
||||||
|
}; |
||||||
|
|
||||||
|
//
|
||||||
|
// InterceptRecvTrailingMetadataLoadBalancingPolicy
|
||||||
|
//
|
||||||
|
|
||||||
|
constexpr char kInterceptRecvTrailingMetadataLbPolicyName[] = |
||||||
|
"intercept_trailing_metadata_lb"; |
||||||
|
|
||||||
|
class InterceptRecvTrailingMetadataLoadBalancingPolicy |
||||||
|
: public ForwardingLoadBalancingPolicy { |
||||||
|
public: |
||||||
|
InterceptRecvTrailingMetadataLoadBalancingPolicy( |
||||||
|
const Args& args, InterceptRecvTrailingMetadataCallback cb, |
||||||
|
void* user_data) |
||||||
|
: ForwardingLoadBalancingPolicy(args, |
||||||
|
/*delegate_lb_policy_name=*/"pick_first"), |
||||||
|
cb_(cb), |
||||||
|
user_data_(user_data) {} |
||||||
|
|
||||||
|
~InterceptRecvTrailingMetadataLoadBalancingPolicy() override = default; |
||||||
|
|
||||||
|
const char* name() const override { |
||||||
|
return kInterceptRecvTrailingMetadataLbPolicyName; |
||||||
|
} |
||||||
|
|
||||||
|
bool PickLocked(PickState* pick, grpc_error** error) override { |
||||||
|
bool ret = ForwardingLoadBalancingPolicy::PickLocked(pick, error); |
||||||
|
// Note: This assumes that the delegate policy does not
|
||||||
|
// intercepting recv_trailing_metadata. If we ever need to use
|
||||||
|
// this with a delegate policy that does, then we'll need to
|
||||||
|
// handle async pick returns separately.
|
||||||
|
New<TrailingMetadataHandler>(pick, cb_, user_data_); // deletes itself
|
||||||
|
return ret; |
||||||
|
} |
||||||
|
|
||||||
|
private: |
||||||
|
class TrailingMetadataHandler { |
||||||
|
public: |
||||||
|
TrailingMetadataHandler(PickState* pick, |
||||||
|
InterceptRecvTrailingMetadataCallback cb, |
||||||
|
void* user_data) |
||||||
|
: cb_(cb), user_data_(user_data) { |
||||||
|
GRPC_CLOSURE_INIT(&recv_trailing_metadata_ready_, |
||||||
|
RecordRecvTrailingMetadata, this, |
||||||
|
grpc_schedule_on_exec_ctx); |
||||||
|
pick->recv_trailing_metadata_ready = &recv_trailing_metadata_ready_; |
||||||
|
pick->original_recv_trailing_metadata_ready = |
||||||
|
&original_recv_trailing_metadata_ready_; |
||||||
|
pick->recv_trailing_metadata = &recv_trailing_metadata_; |
||||||
|
} |
||||||
|
|
||||||
|
private: |
||||||
|
static void RecordRecvTrailingMetadata(void* arg, grpc_error* err) { |
||||||
|
TrailingMetadataHandler* self = |
||||||
|
static_cast<TrailingMetadataHandler*>(arg); |
||||||
|
GPR_ASSERT(self->recv_trailing_metadata_ != nullptr); |
||||||
|
self->cb_(self->user_data_); |
||||||
|
GRPC_CLOSURE_SCHED(self->original_recv_trailing_metadata_ready_, |
||||||
|
GRPC_ERROR_REF(err)); |
||||||
|
Delete(self); |
||||||
|
} |
||||||
|
|
||||||
|
InterceptRecvTrailingMetadataCallback cb_; |
||||||
|
void* user_data_; |
||||||
|
grpc_closure recv_trailing_metadata_ready_; |
||||||
|
grpc_closure* original_recv_trailing_metadata_ready_ = nullptr; |
||||||
|
grpc_metadata_batch* recv_trailing_metadata_ = nullptr; |
||||||
|
}; |
||||||
|
|
||||||
|
InterceptRecvTrailingMetadataCallback cb_; |
||||||
|
void* user_data_; |
||||||
|
}; |
||||||
|
|
||||||
|
class InterceptTrailingFactory : public LoadBalancingPolicyFactory { |
||||||
|
public: |
||||||
|
explicit InterceptTrailingFactory(InterceptRecvTrailingMetadataCallback cb, |
||||||
|
void* user_data) |
||||||
|
: cb_(cb), user_data_(user_data) {} |
||||||
|
|
||||||
|
grpc_core::OrphanablePtr<grpc_core::LoadBalancingPolicy> |
||||||
|
CreateLoadBalancingPolicy( |
||||||
|
const grpc_core::LoadBalancingPolicy::Args& args) const override { |
||||||
|
return grpc_core::OrphanablePtr<grpc_core::LoadBalancingPolicy>( |
||||||
|
grpc_core::New<InterceptRecvTrailingMetadataLoadBalancingPolicy>( |
||||||
|
args, cb_, user_data_)); |
||||||
|
} |
||||||
|
|
||||||
|
const char* name() const override { |
||||||
|
return kInterceptRecvTrailingMetadataLbPolicyName; |
||||||
|
} |
||||||
|
|
||||||
|
private: |
||||||
|
InterceptRecvTrailingMetadataCallback cb_; |
||||||
|
void* user_data_; |
||||||
|
}; |
||||||
|
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
void RegisterInterceptRecvTrailingMetadataLoadBalancingPolicy( |
||||||
|
InterceptRecvTrailingMetadataCallback cb, void* user_data) { |
||||||
|
grpc_core::LoadBalancingPolicyRegistry::Builder:: |
||||||
|
RegisterLoadBalancingPolicyFactory( |
||||||
|
grpc_core::UniquePtr<grpc_core::LoadBalancingPolicyFactory>( |
||||||
|
grpc_core::New<InterceptTrailingFactory>(cb, user_data))); |
||||||
|
} |
||||||
|
|
||||||
|
} // namespace grpc_core
|
@ -0,0 +1,34 @@ |
|||||||
|
/*
|
||||||
|
* |
||||||
|
* Copyright 2018 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. |
||||||
|
* |
||||||
|
*/ |
||||||
|
|
||||||
|
#ifndef GRPC_TEST_CORE_UTIL_TEST_LB_POLICIES_H |
||||||
|
#define GRPC_TEST_CORE_UTIL_TEST_LB_POLICIES_H |
||||||
|
|
||||||
|
namespace grpc_core { |
||||||
|
|
||||||
|
typedef void (*InterceptRecvTrailingMetadataCallback)(void*); |
||||||
|
|
||||||
|
// Registers an LB policy called "intercept_trailing_metadata_lb" that
|
||||||
|
// invokes cb with argument user_data when trailing metadata is received
|
||||||
|
// for each call.
|
||||||
|
void RegisterInterceptRecvTrailingMetadataLoadBalancingPolicy( |
||||||
|
InterceptRecvTrailingMetadataCallback cb, void* user_data); |
||||||
|
|
||||||
|
} // namespace grpc_core
|
||||||
|
|
||||||
|
#endif // GRPC_TEST_CORE_UTIL_TEST_LB_POLICIES_H
|
Loading…
Reference in new issue