%% ------------------------------------------------------------------- %% %% Copyright (c) 2014 Basho Technologies, Inc. All Rights Reserved. %% Copyright (c) 2016 Christopher Meiklejohn. All Rights Reserved. %% %% This file is provided 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(partisan_util). -include("partisan_logger.hrl"). -include("partisan.hrl"). -export([maps_append/3]). -export([encode/1]). -export([encode/2]). -export([maybe_connect_disterl/1]). -export([maybe_pad_term/1]). -export([get/2]). -export([get/3]). %% ============================================================================= %% API %% ============================================================================= %% ----------------------------------------------------------------------------- %% @doc %% @end %% ----------------------------------------------------------------------------- -spec get(Key :: term(), Arg :: list() | map()) -> Value :: any() | no_return(). get(Key, Arg) when is_list(Arg) -> case lists:keyfind(Key, 1, Arg) of {Key, Value} -> Value; false -> error(badkey) end; get(Key, Arg) when is_map(Arg) -> maps:get(Key, Arg). %% ----------------------------------------------------------------------------- %% @doc %% @end %% ----------------------------------------------------------------------------- -spec get(Key :: term(), Arg :: list() | map(), Default :: any()) -> Value :: any(). get(Key, Arg, Default) when is_list(Arg) -> case lists:keyfind(Key, 1, Arg) of {Key, Value} -> Value; false -> Default end; get(Key, Arg, Default) when is_map(Arg) -> maps:get(Key, Arg, Default). %% ----------------------------------------------------------------------------- %% @doc %% @end %% ----------------------------------------------------------------------------- -spec maybe_connect_disterl(Node :: node()) -> ok. maybe_connect_disterl(Node ) -> case partisan_config:get(connect_disterl, false) of true -> _ = net_kernel:connect_node(Node), ok; false -> ok end. %% ----------------------------------------------------------------------------- %% @doc %% @end %% ----------------------------------------------------------------------------- -spec maybe_pad_term(Term :: term()) -> term() | {'$partisan_padded', Padding :: term(), Term :: term()}. maybe_pad_term(Term) -> case partisan_config:get(binary_padding, false) of true -> Padding = partisan_config:get(binary_padding_term, undefined), {'$partisan_padded', Padding, Term}; false -> Term end. %% ----------------------------------------------------------------------------- %% @doc %% @end %% ----------------------------------------------------------------------------- maps_append(Key, Value, Map) -> maps:update_with( Key, fun ([]) -> [Value]; (L) when is_list(L) -> [Value | L]; (X) -> [Value, X] end, [Value], Map ). %% ----------------------------------------------------------------------------- %% @doc %% @end %% ----------------------------------------------------------------------------- encode(Term) -> encode(Term, []). %% ----------------------------------------------------------------------------- %% @doc If `pid_encoding' or `ref_encoding' configuration options are enabled, %% this function will ignore Opts and return an iolist(). Otherwise, the %% function calls `erlang:term_to_iovec(Term, Opts)'.' %% @end %% ----------------------------------------------------------------------------- -spec encode(Term :: term(), Opts :: list()) -> erlang:ext_iovec() | iolist(). encode(Term, Opts) -> PidEnc = partisan_config:get(pid_encoding, true), RefEnc = partisan_config:get(ref_encoding, true), case PidEnc orelse RefEnc of true -> erlang:iolist_to_iovec([131, encode_(Term)]); false -> erlang:term_to_iovec(Term, Opts) end. %% ============================================================================= %% PRIVATE %% ============================================================================= %% ----------------------------------------------------------------------------- %% @doc %% @end %% ----------------------------------------------------------------------------- encode_([]) -> 106; encode_({}) -> [104, 0]; encode_(T) when is_atom(T) -> <<131, Rest/binary>> = term_to_binary(T), [Rest]; encode_(T) when is_binary(T) -> Len = byte_size(T), [109, <>, T]; encode_(T) when is_tuple(T) -> Len = tuple_size(T), case Len > 255 of false -> [104, Len, [encode_(E) || E <- tuple_to_list(T)]]; true -> [104, <>, [encode_(E) || E <- tuple_to_list(T)]] end; encode_([alias|Ref]) -> %% Improper list with 2 elements to support partisan_remote_ref format [108, <<1:32/integer-big>>, [encode_(alias), encode_(Ref)]]; encode_([Node|Bin]) when is_atom(Node), is_binary(Bin) -> %% Improper list with 2 elements to support partisan_remote_ref format [108, <<1:32/integer-big>>, [encode_(Node), encode_(Bin)]]; encode_(T) when is_list(T) -> %% TODO improper lists Len = length(T), case Len < 64436 andalso lists:all(fun(E) when is_integer(E), E >= 0, E < 256 -> true; (_) -> false end, T) of true -> [107, <>, T]; false -> [108, <>, [[encode_(E) || E <- T]], 106] end; encode_(T) when is_map(T) -> Len = maps:size(T), [116, <>, [[encode_(K), encode_(V)] || {K, V} <- maps:to_list(T)]]; encode_(T) when is_reference(T) -> case partisan_config:get(ref_encoding, true) of false -> <<131, Rest/binary>> = term_to_binary(T), Rest; true -> <<131, Rest/binary>> = term_to_binary( partisan_remote_ref:from_term(T) ), Rest end; encode_(T) when is_pid(T) -> case partisan_config:get(pid_encoding, true) of false -> <<131, Rest/binary>> = term_to_binary(T), Rest; true -> <<131, Rest/binary>> = term_to_binary(pid(T)), Rest end; encode_(T) -> %% fallback clause <<131, Rest/binary>> = term_to_binary(T), Rest. %% ----------------------------------------------------------------------------- %% @private %% @doc %% @end %% ----------------------------------------------------------------------------- -spec pid(pid()) -> partisan:remote_pid(). pid(Pid) when is_pid(Pid) -> Node = erlang:node(Pid), Disterl = partisan_config:get(connect_disterl, false), pid(Pid, Node, Disterl). %% @private pid(Pid, _, false) -> partisan_remote_ref:from_term(Pid); pid(Pid, Node, true) -> case Node == node() of true -> %% This is super dangerous. case partisan_config:get(register_pid_for_encoding, false) of true -> Name = case process_info(Pid, registered_name) of {registered_name, OldName} -> ?LOG_DEBUG( "unregistering pid: ~p with name: ~p", [Pid, OldName] ), %% TODO: Race condition on unregister/register. unregister(OldName), OldName; [] -> Unique = erlang:unique_integer( [monotonic, positive] ), list_to_atom( "partisan_registered_name_" ++ integer_to_list(Unique) ) end, ?LOG_DEBUG( "registering pid: ~p as name: ~p at node: ~p", [Pid, Name, Node] ), true = erlang:register(Name, Pid), partisan_remote_ref:from_term(Name, Node); false -> partisan_remote_ref:from_term(Pid, Node) end; false -> %% This is even more super dangerous. case partisan_config:get(register_pid_for_encoding, false) of true -> Unique = erlang:unique_integer([monotonic, positive]), Name = list_to_atom( "partisan_registered_name_" ++ integer_to_list(Unique) ), [_, B, C] = string:split(pid_to_list(Pid), ".", all), PidString = "<0." ++ B ++ "." ++ C, ?LOG_DEBUG( "registering pid: ~p as name: ~p at node: ~p", [Pid, Name, Node] ), RegisterFun = fun() -> LocalPid = list_to_pid(PidString), RegName = process_info(LocalPid, registered_name), case RegName of {registered_name, OldName} when is_atom(OldName) -> ?LOG_DEBUG( "unregistering pid: ~p with name: ~p", [Pid, OldName] ), %% TODO: Race condition on unregister/ %% register. unregister(OldName), ok; _ -> ok end, erlang:register(Name, LocalPid) end, %% TODO: Race here unless we wait. _ = partisan_rpc:call( Node, erlang, spawn, [RegisterFun], 5000 ), partisan_remote_ref:from_term(Name, Node); false -> [_, B, C] = string:split(pid_to_list(Pid), ".", all), PidString = "<0." ++ B ++ "." ++ C, partisan_remote_ref:from_term(list_to_pid(PidString), Node) end end.