| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you 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. |
| |
| use log::debug; |
| |
| use std::collections::HashMap; |
| use std::convert::Into; |
| use std::fmt; |
| use std::fmt::{Debug, Formatter}; |
| use std::sync::{Arc, Mutex}; |
| |
| use crate::protocol::{TInputProtocol, TMessageIdentifier, TOutputProtocol, TStoredInputProtocol}; |
| |
| use super::{handle_process_result, TProcessor}; |
| |
| const MISSING_SEPARATOR_AND_NO_DEFAULT: &str = |
| "missing service separator and no default processor set"; |
| type ThreadSafeProcessor = Box<dyn TProcessor + Send + Sync>; |
| |
| /// A `TProcessor` that can demux service calls to multiple underlying |
| /// Thrift services. |
| /// |
| /// Users register service-specific `TProcessor` instances with a |
| /// `TMultiplexedProcessor`, and then register that processor with a server |
| /// implementation. Following that, all incoming service calls are automatically |
| /// routed to the service-specific `TProcessor`. |
| /// |
| /// A `TMultiplexedProcessor` can only handle messages sent by a |
| /// `TMultiplexedOutputProtocol`. |
| #[derive(Default)] |
| pub struct TMultiplexedProcessor { |
| stored: Mutex<StoredProcessors>, |
| } |
| |
| #[derive(Default)] |
| struct StoredProcessors { |
| processors: HashMap<String, Arc<ThreadSafeProcessor>>, |
| default_processor: Option<Arc<ThreadSafeProcessor>>, |
| } |
| |
| impl TMultiplexedProcessor { |
| /// Create a new `TMultiplexedProcessor` with no registered service-specific |
| /// processors. |
| pub fn new() -> TMultiplexedProcessor { |
| TMultiplexedProcessor { |
| stored: Mutex::new(StoredProcessors { |
| processors: HashMap::new(), |
| default_processor: None, |
| }), |
| } |
| } |
| |
| /// Register a service-specific `processor` for the service named |
| /// `service_name`. This implementation is also backwards-compatible with |
| /// non-multiplexed clients. Set `as_default` to `true` to allow |
| /// non-namespaced requests to be dispatched to a default processor. |
| /// |
| /// Returns success if a new entry was inserted. Returns an error if: |
| /// * A processor exists for `service_name` |
| /// * You attempt to register a processor as default, and an existing default exists |
| #[allow(clippy::map_entry)] |
| pub fn register<S: Into<String>>( |
| &mut self, |
| service_name: S, |
| processor: Box<dyn TProcessor + Send + Sync>, |
| as_default: bool, |
| ) -> crate::Result<()> { |
| let mut stored = self.stored.lock().unwrap(); |
| |
| let name = service_name.into(); |
| if !stored.processors.contains_key(&name) { |
| let processor = Arc::new(processor); |
| |
| if as_default { |
| if stored.default_processor.is_none() { |
| stored.processors.insert(name, processor.clone()); |
| stored.default_processor = Some(processor.clone()); |
| Ok(()) |
| } else { |
| Err("cannot reset default processor".into()) |
| } |
| } else { |
| stored.processors.insert(name, processor); |
| Ok(()) |
| } |
| } else { |
| Err(format!("cannot overwrite existing processor for service {}", name).into()) |
| } |
| } |
| |
| fn process_message( |
| &self, |
| msg_ident: &TMessageIdentifier, |
| i_prot: &mut dyn TInputProtocol, |
| o_prot: &mut dyn TOutputProtocol, |
| ) -> crate::Result<()> { |
| let (svc_name, svc_call) = split_ident_name(&msg_ident.name); |
| debug!("routing svc_name {:?} svc_call {}", &svc_name, &svc_call); |
| |
| let processor: Option<Arc<ThreadSafeProcessor>> = { |
| let stored = self.stored.lock().unwrap(); |
| if let Some(name) = svc_name { |
| stored.processors.get(name).cloned() |
| } else { |
| stored.default_processor.clone() |
| } |
| }; |
| |
| match processor { |
| Some(arc) => { |
| let new_msg_ident = TMessageIdentifier::new( |
| svc_call, |
| msg_ident.message_type, |
| msg_ident.sequence_number, |
| ); |
| let mut proxy_i_prot = TStoredInputProtocol::new(i_prot, new_msg_ident); |
| (*arc).process(&mut proxy_i_prot, o_prot) |
| } |
| None => Err(missing_processor_message(svc_name).into()), |
| } |
| } |
| } |
| |
| impl TProcessor for TMultiplexedProcessor { |
| fn process(&self, i_prot: &mut dyn TInputProtocol, o_prot: &mut dyn TOutputProtocol) -> crate::Result<()> { |
| let msg_ident = i_prot.read_message_begin()?; |
| |
| debug!("process incoming msg id:{:?}", &msg_ident); |
| let res = self.process_message(&msg_ident, i_prot, o_prot); |
| |
| handle_process_result(&msg_ident, res, o_prot) |
| } |
| } |
| |
| impl Debug for TMultiplexedProcessor { |
| fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { |
| let stored = self.stored.lock().unwrap(); |
| write!( |
| f, |
| "TMultiplexedProcess {{ registered_count: {:?} default: {:?} }}", |
| stored.processors.keys().len(), |
| stored.default_processor.is_some() |
| ) |
| } |
| } |
| |
| fn split_ident_name(ident_name: &str) -> (Option<&str>, &str) { |
| ident_name |
| .find(':') |
| .map(|pos| { |
| let (svc_name, svc_call) = ident_name.split_at(pos); |
| let (_, svc_call) = svc_call.split_at(1); // remove colon from service call name |
| (Some(svc_name), svc_call) |
| }) |
| .or_else(|| Some((None, ident_name))) |
| .unwrap() |
| } |
| |
| fn missing_processor_message(svc_name: Option<&str>) -> String { |
| match svc_name { |
| Some(name) => format!("no processor found for service {}", name), |
| None => MISSING_SEPARATOR_AND_NO_DEFAULT.to_owned(), |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use std::convert::Into; |
| use std::sync::atomic::{AtomicBool, Ordering}; |
| use std::sync::Arc; |
| |
| use crate::protocol::{TBinaryInputProtocol, TBinaryOutputProtocol, TMessageIdentifier, TMessageType}; |
| use crate::transport::{ReadHalf, TBufferChannel, TIoChannel, WriteHalf}; |
| use crate::{ApplicationError, ApplicationErrorKind}; |
| |
| use super::*; |
| |
| #[test] |
| fn should_split_name_into_proper_separator_and_service_call() { |
| let ident_name = "foo:bar_call"; |
| let (serv, call) = split_ident_name(&ident_name); |
| assert_eq!(serv, Some("foo")); |
| assert_eq!(call, "bar_call"); |
| } |
| |
| #[test] |
| fn should_return_full_ident_if_no_separator_exists() { |
| let ident_name = "bar_call"; |
| let (serv, call) = split_ident_name(&ident_name); |
| assert_eq!(serv, None); |
| assert_eq!(call, "bar_call"); |
| } |
| |
| #[test] |
| fn should_write_error_if_no_separator_found_and_no_default_processor_exists() { |
| let (mut i, mut o) = build_objects(); |
| |
| let sent_ident = TMessageIdentifier::new("foo", TMessageType::Call, 10); |
| o.write_message_begin(&sent_ident).unwrap(); |
| o.flush().unwrap(); |
| o.transport.copy_write_buffer_to_read_buffer(); |
| o.transport.empty_write_buffer(); |
| |
| let p = TMultiplexedProcessor::new(); |
| p.process(&mut i, &mut o).unwrap(); // at this point an error should be written out |
| |
| i.transport.set_readable_bytes(&o.transport.write_bytes()); |
| let rcvd_ident = i.read_message_begin().unwrap(); |
| let expected_ident = TMessageIdentifier::new("foo", TMessageType::Exception, 10); |
| assert_eq!(rcvd_ident, expected_ident); |
| let rcvd_err = crate::Error::read_application_error_from_in_protocol(&mut i).unwrap(); |
| let expected_err = ApplicationError::new( |
| ApplicationErrorKind::Unknown, |
| MISSING_SEPARATOR_AND_NO_DEFAULT, |
| ); |
| assert_eq!(rcvd_err, expected_err); |
| } |
| |
| #[test] |
| fn should_write_error_if_separator_exists_and_no_processor_found() { |
| let (mut i, mut o) = build_objects(); |
| |
| let sent_ident = TMessageIdentifier::new("missing:call", TMessageType::Call, 10); |
| o.write_message_begin(&sent_ident).unwrap(); |
| o.flush().unwrap(); |
| o.transport.copy_write_buffer_to_read_buffer(); |
| o.transport.empty_write_buffer(); |
| |
| let p = TMultiplexedProcessor::new(); |
| p.process(&mut i, &mut o).unwrap(); // at this point an error should be written out |
| |
| i.transport.set_readable_bytes(&o.transport.write_bytes()); |
| let rcvd_ident = i.read_message_begin().unwrap(); |
| let expected_ident = TMessageIdentifier::new("missing:call", TMessageType::Exception, 10); |
| assert_eq!(rcvd_ident, expected_ident); |
| let rcvd_err = crate::Error::read_application_error_from_in_protocol(&mut i).unwrap(); |
| let expected_err = ApplicationError::new( |
| ApplicationErrorKind::Unknown, |
| missing_processor_message(Some("missing")), |
| ); |
| assert_eq!(rcvd_err, expected_err); |
| } |
| |
| #[derive(Default)] |
| struct Service { |
| pub invoked: Arc<AtomicBool>, |
| } |
| |
| impl TProcessor for Service { |
| fn process(&self, _: &mut dyn TInputProtocol, _: &mut dyn TOutputProtocol) -> crate::Result<()> { |
| let res = self |
| .invoked |
| .compare_and_swap(false, true, Ordering::Relaxed); |
| if res { |
| Ok(()) |
| } else { |
| Err("failed swap".into()) |
| } |
| } |
| } |
| |
| #[test] |
| fn should_route_call_to_correct_processor() { |
| let (mut i, mut o) = build_objects(); |
| |
| // build the services |
| let svc_1 = Service { |
| invoked: Arc::new(AtomicBool::new(false)), |
| }; |
| let atm_1 = svc_1.invoked.clone(); |
| let svc_2 = Service { |
| invoked: Arc::new(AtomicBool::new(false)), |
| }; |
| let atm_2 = svc_2.invoked.clone(); |
| |
| // register them |
| let mut p = TMultiplexedProcessor::new(); |
| p.register("service_1", Box::new(svc_1), false).unwrap(); |
| p.register("service_2", Box::new(svc_2), false).unwrap(); |
| |
| // make the service call |
| let sent_ident = TMessageIdentifier::new("service_1:call", TMessageType::Call, 10); |
| o.write_message_begin(&sent_ident).unwrap(); |
| o.flush().unwrap(); |
| o.transport.copy_write_buffer_to_read_buffer(); |
| o.transport.empty_write_buffer(); |
| |
| p.process(&mut i, &mut o).unwrap(); |
| |
| // service 1 should have been invoked, not service 2 |
| assert_eq!(atm_1.load(Ordering::Relaxed), true); |
| assert_eq!(atm_2.load(Ordering::Relaxed), false); |
| } |
| |
| #[test] |
| fn should_route_call_to_correct_processor_if_no_separator_exists_and_default_processor_set() { |
| let (mut i, mut o) = build_objects(); |
| |
| // build the services |
| let svc_1 = Service { |
| invoked: Arc::new(AtomicBool::new(false)), |
| }; |
| let atm_1 = svc_1.invoked.clone(); |
| let svc_2 = Service { |
| invoked: Arc::new(AtomicBool::new(false)), |
| }; |
| let atm_2 = svc_2.invoked.clone(); |
| |
| // register them |
| let mut p = TMultiplexedProcessor::new(); |
| p.register("service_1", Box::new(svc_1), false).unwrap(); |
| p.register("service_2", Box::new(svc_2), true).unwrap(); // second processor is default |
| |
| // make the service call (it's an old client, so we have to be backwards compatible) |
| let sent_ident = TMessageIdentifier::new("old_call", TMessageType::Call, 10); |
| o.write_message_begin(&sent_ident).unwrap(); |
| o.flush().unwrap(); |
| o.transport.copy_write_buffer_to_read_buffer(); |
| o.transport.empty_write_buffer(); |
| |
| p.process(&mut i, &mut o).unwrap(); |
| |
| // service 2 should have been invoked, not service 1 |
| assert_eq!(atm_1.load(Ordering::Relaxed), false); |
| assert_eq!(atm_2.load(Ordering::Relaxed), true); |
| } |
| |
| fn build_objects() -> ( |
| TBinaryInputProtocol<ReadHalf<TBufferChannel>>, |
| TBinaryOutputProtocol<WriteHalf<TBufferChannel>>, |
| ) { |
| let c = TBufferChannel::with_capacity(128, 128); |
| let (r_c, w_c) = c.split().unwrap(); |
| ( |
| TBinaryInputProtocol::new(r_c, true), |
| TBinaryOutputProtocol::new(w_c, true), |
| ) |
| } |
| } |