lib/erl/src/thrift_processor.erl

%%
%% 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.
%%

-module(thrift_processor).

-export([init/1]).

-include("thrift_constants.hrl").
-include("thrift_protocol.hrl").

-record(thrift_processor, {handler, protocol, service}).

init({_Server, ProtoGen, Service, Handler}) when is_function(ProtoGen, 0) ->
    {ok, Proto} = ProtoGen(),
    loop(#thrift_processor{protocol = Proto,
                           service = Service,
                           handler = Handler}).

loop(State0 = #thrift_processor{protocol  = Proto0,
                                handler = Handler,
                                service = Service}) ->

    {Proto1, MessageBegin} = thrift_protocol:read(Proto0, message_begin),
    State1 = State0#thrift_processor{protocol = Proto1},

    ErrorHandler = fun
        (HandlerModules) when is_list(HandlerModules) -> thrift_multiplexed_map_wrapper:fetch(?MULTIPLEXED_ERROR_HANDLER_KEY, HandlerModules);
        (HandlerModule) -> HandlerModule
    end,

    case MessageBegin of

        #protocol_message_begin{name = Function,
                                type = Type,
                                seqid = Seqid} when Type =:= ?tMessageType_CALL; Type =:= ?tMessageType_ONEWAY ->
            case string:tokens(Function, ?MULTIPLEXED_SERVICE_SEPARATOR) of
                [ServiceName, FunctionName] ->
                    ServiceModule  = thrift_multiplexed_map_wrapper:fetch(ServiceName, Service),
                    ServiceHandler = thrift_multiplexed_map_wrapper:fetch(ServiceName, Handler),
                    case handle_function(State1#thrift_processor{service=ServiceModule, handler=ServiceHandler}, list_to_atom(FunctionName), Seqid) of
                        {State2, ok} -> loop(State2#thrift_processor{service=Service, handler=Handler});
                        {_State2, {error, Reason}} ->
							apply(ErrorHandler(Handler), handle_error, [list_to_atom(Function), Reason]),
                            thrift_protocol:close_transport(Proto1),
                            ok
                    end;
                _ ->
                    case handle_function(State1, list_to_atom(Function), Seqid) of
                        {State2, ok} -> loop(State2);
                        {_State2, {error, Reason}} ->
							apply(ErrorHandler(Handler), handle_error, [list_to_atom(Function), Reason]),
                            thrift_protocol:close_transport(Proto1),
                            ok
                    end
            end;
        {error, timeout = Reason} ->
			apply(ErrorHandler(Handler), handle_error, [undefined, Reason]),
            thrift_protocol:close_transport(Proto1),
            ok;
        {error, closed = Reason} ->
            %% error_logger:info_msg("Client disconnected~n"),
			apply(ErrorHandler(Handler), handle_error, [undefined, Reason]),
            thrift_protocol:close_transport(Proto1),
            exit(shutdown);
        {error, Reason} ->
			apply(ErrorHandler(Handler), handle_error, [undefined, Reason]),
            thrift_protocol:close_transport(Proto1),
            exit(shutdown)
    end.

handle_function(State0=#thrift_processor{protocol = Proto0,
                                         handler = Handler,
                                         service = Service},
                Function,
                Seqid) ->
    InParams = Service:function_info(Function, params_type),

    {Proto1, {ok, Params}} = thrift_protocol:read(Proto0, InParams),
    State1 = State0#thrift_processor{protocol = Proto1},

    try
        Result = Handler:handle_function(Function, Params),
        %% {Micro, Result} = better_timer(Handler, handle_function, [Function, Params]),
        %% error_logger:info_msg("Processed ~p(~p) in ~.4fms~n",
        %%                       [Function, Params, Micro/1000.0]),
        handle_success(State1, Function, Result, Seqid)
    catch
        Type:Data when Type =:= throw orelse Type =:= error ->
            handle_function_catch(State1, Function, Type, Data, Seqid)
    end.

handle_function_catch(State = #thrift_processor{service = Service},
                      Function, ErrType, ErrData, Seqid) ->
    IsOneway = Service:function_info(Function, reply_type) =:= oneway_void,

    case {ErrType, ErrData} of
        _ when IsOneway ->
            Stack = erlang:get_stacktrace(),
            error_logger:warning_msg(
              "oneway void ~p threw error which must be ignored: ~p",
              [Function, {ErrType, ErrData, Stack}]),
            {State, ok};

        {throw, Exception} when is_tuple(Exception), size(Exception) > 0 ->
            %error_logger:warning_msg("~p threw exception: ~p~n", [Function, Exception]),
            handle_exception(State, Function, Exception, Seqid);
            % we still want to accept more requests from this client

        {error, Error} ->
            handle_error(State, Function, Error, Seqid)
    end.

handle_success(State = #thrift_processor{service = Service},
               Function,
               Result,
               Seqid) ->
    ReplyType  = Service:function_info(Function, reply_type),
    StructName = atom_to_list(Function) ++ "_result",

    case Result of
        {reply, ReplyData} ->
            Reply = {{struct, [{0, ReplyType}]}, {StructName, ReplyData}},
            send_reply(State, Function, ?tMessageType_REPLY, Reply, Seqid);

        ok when ReplyType == {struct, []} ->
            send_reply(State, Function, ?tMessageType_REPLY, {ReplyType, {StructName}}, Seqid);

        ok when ReplyType == oneway_void ->
            %% no reply for oneway void
            {State, ok}
    end.

handle_exception(State = #thrift_processor{service = Service},
                 Function,
                 Exception,
                 Seqid) ->
    ExceptionType = element(1, Exception),
    %% Fetch a structure like {struct, [{-2, {struct, {Module, Type}}},
    %%                                  {-3, {struct, {Module, Type}}}]}

    ReplySpec = Service:function_info(Function, exceptions),
    {struct, XInfo} = ReplySpec,

    true = is_list(XInfo),

    %% Assuming we had a type1 exception, we'd get: [undefined, Exception, undefined]
    %% e.g.: [{-1, type0}, {-2, type1}, {-3, type2}]
    ExceptionList = [case Type of
                         ExceptionType -> Exception;
                         _ -> undefined
                     end
                     || {_Fid, {struct, {_Module, Type}}} <- XInfo],

    ExceptionTuple = list_to_tuple([Function | ExceptionList]),

                                                % Make sure we got at least one defined
    case lists:all(fun(X) -> X =:= undefined end, ExceptionList) of
        true ->
            handle_unknown_exception(State, Function, Exception, Seqid);
        false ->
            send_reply(State, Function, ?tMessageType_REPLY, {ReplySpec, ExceptionTuple}, Seqid)
    end.

%%
%% Called when an exception has been explicitly thrown by the service, but it was
%% not one of the exceptions that was defined for the function.
%%
handle_unknown_exception(State, Function, Exception, Seqid) ->
    handle_error(State, Function, {exception_not_declared_as_thrown,
                                   Exception}, Seqid).

handle_error(State, Function, Error, Seqid) ->
    Stack = erlang:get_stacktrace(),
    error_logger:error_msg("~p had an error: ~p~n", [Function, {Error, Stack}]),

    Message =
        case application:get_env(thrift, exceptions_include_traces) of
            {ok, true} ->
                lists:flatten(io_lib:format("An error occurred: ~p~n",
                                            [{Error, Stack}]));
            _ ->
                "An unknown handler error occurred."
        end,
    Reply = {?TApplicationException_Structure,
             #'TApplicationException'{
                message = Message,
                type = ?TApplicationException_UNKNOWN}},
    send_reply(State, Function, ?tMessageType_EXCEPTION, Reply, Seqid).

send_reply(State = #thrift_processor{protocol = Proto0}, Function, ReplyMessageType, Reply, Seqid) ->
    try
        {Proto1, ok} = thrift_protocol:write(Proto0, #protocol_message_begin{
                                               name = atom_to_list(Function),
                                               type = ReplyMessageType,
                                               seqid = Seqid}),
        {Proto2, ok} = thrift_protocol:write(Proto1, Reply),
        {Proto3, ok} = thrift_protocol:write(Proto2, message_end),
        {Proto4, ok} = thrift_protocol:flush_transport(Proto3),
        {State#thrift_processor{protocol = Proto4}, ok}
    catch
        error:{badmatch, {_, {error, _} = Error}} ->
            {State, Error}
    end.