├── .gitignore ├── rebar ├── include └── riak_exchange.hrl ├── src ├── rabbit_exchange_type_riak.app.src └── rabbit_exchange_type_riak.erl ├── priv ├── send_lots_of_msgs.py ├── subscriber.py └── send_msg.py ├── rebar.config ├── Makefile └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | deps 3 | ebin 4 | dist 5 | .idea 6 | *.i* -------------------------------------------------------------------------------- /rebar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jbrisbin/riak-exchange/HEAD/rebar -------------------------------------------------------------------------------- /include/riak_exchange.hrl: -------------------------------------------------------------------------------- 1 | -include_lib("rabbit_common/include/rabbit.hrl"). 2 | -include_lib("rabbit_common/include/rabbit_framing.hrl"). 3 | -------------------------------------------------------------------------------- /src/rabbit_exchange_type_riak.app.src: -------------------------------------------------------------------------------- 1 | {application, rabbit_exchange_type_riak, 2 | [ 3 | {description, "RabbitMQ Riak Exchange Plugin"}, 4 | {vsn, "0.2.0"}, 5 | {modules, []}, 6 | {registered, []}, 7 | {env, []}, 8 | {applications, [ 9 | kernel, 10 | stdlib, 11 | rabbit, 12 | mnesia, 13 | protobuffs, 14 | riakc 15 | ]} 16 | ] 17 | }. 18 | -------------------------------------------------------------------------------- /priv/send_lots_of_msgs.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import time 3 | import amqplib.client_0_8 as amqp 4 | 5 | conn = amqp.Connection() 6 | ch = conn.channel() 7 | 8 | start = time.time() 9 | 10 | for i in range(10000): 11 | msg = amqp.Message( 12 | "Message #%s" % i, 13 | content_type="text/plain", 14 | application_headers={ 15 | "test": "header" 16 | } 17 | ) 18 | ch.basic_publish(msg, "riak", "msg.%s" % i) 19 | 20 | end = time.time() 21 | print "Elapsed time: %s" % (10000 / (end - start)) -------------------------------------------------------------------------------- /rebar.config: -------------------------------------------------------------------------------- 1 | { 2 | deps, [ 3 | {rabbit_common, ".*", {git, "https://github.com/jbrisbin/rabbit_common.git", {tag, "rabbitmq-3.3.5"}}}, 4 | {riakc, ".*", {git, "https://github.com/basho/riak-erlang-client.git", {tag, "master"}}} 5 | ] 6 | }. 7 | 8 | {erl_opts, [ 9 | debug_info, 10 | compressed, 11 | report, 12 | warn_export_all, 13 | warn_export_vars, 14 | warn_shadow_vars, 15 | warn_unused_function, 16 | warn_deprecated_function, 17 | warn_obsolete_guard, 18 | warn_unused_import 19 | % warnings_as_errors 20 | ]}. 21 | -------------------------------------------------------------------------------- /priv/subscriber.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import amqplib.client_0_8 as amqp 3 | 4 | conn = amqp.Connection() 5 | ch = conn.channel() 6 | 7 | def callback(msg): 8 | print "msg: %s" % msg.body 9 | msg.channel.basic_ack(msg.delivery_tag) 10 | msg.channel.basic_cancel(msg.consumer_tag) 11 | 12 | ch.access_request('/rtest', active=True, read=True) 13 | 14 | qname, _, _ = ch.queue_declare() 15 | ch.queue_bind(qname, "rtest", "rtest") 16 | ch.basic_consume(qname, callback=callback) 17 | 18 | while ch.callbacks: 19 | ch.wait() 20 | 21 | ch.close() 22 | conn.close() -------------------------------------------------------------------------------- /priv/send_msg.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | from datetime import datetime 3 | import amqplib.client_0_8 as amqp 4 | 5 | conn = amqp.Connection() 6 | ch = conn.channel() 7 | 8 | timestamp = datetime.now().strftime("%y%m%d%H%M%S%f") 9 | msg = amqp.Message( 10 | "{\"type\": \"presence\", \"user\": \"jonbrisbin\", \"timestamp\": \"%s\"}" % (datetime.now().strftime("%d/%m/%y %H:%M")), 11 | content_type="application/json", 12 | application_headers={ 13 | "Location": "Chicago", 14 | "X-Riak-Bucket": "auditlog", 15 | "X-Riak-Key": timestamp 16 | } 17 | ) 18 | ch.basic_publish(msg, "auditlog", timestamp) 19 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | PACKAGE=riak-exchange 2 | DIST_DIR=dist 3 | EBIN_DIR=ebin 4 | INCLUDE_DIRS=include 5 | DEPS_DIR=deps 6 | DEPS ?= riakc protobuffs riak_pb 7 | DEPS_EZ=$(foreach DEP, $(DEPS), $(DEPS_DIR)/$(DEP).ez) 8 | RABBITMQ_HOME ?= . 9 | 10 | all: compile 11 | 12 | clean: 13 | rm -rf $(DIST_DIR) 14 | rm -rf $(EBIN_DIR) 15 | 16 | distclean: clean 17 | rm -rf $(DEPS_DIR) 18 | 19 | package: compile $(DEPS_EZ) 20 | rm -f $(DIST_DIR)/$(PACKAGE).ez 21 | mkdir -p $(DIST_DIR)/$(PACKAGE) 22 | cp -r $(EBIN_DIR) $(DIST_DIR)/$(PACKAGE) 23 | $(foreach EXTRA_DIR, $(INCLUDE_DIRS), cp -r $(EXTRA_DIR) $(DIST_DIR)/$(PACKAGE);) 24 | (cd $(DIST_DIR); zip -r $(PACKAGE).ez $(PACKAGE)) 25 | 26 | install: package 27 | $(foreach DEP, $(DEPS_EZ), cp $(DEP) $(RABBITMQ_HOME)/plugins;) 28 | cp $(DIST_DIR)/$(PACKAGE).ez $(RABBITMQ_HOME)/plugins 29 | 30 | $(DEPS_DIR): 31 | ./rebar get-deps 32 | 33 | $(DEPS_EZ): 34 | cd $(DEPS_DIR); $(foreach DEP, $(DEPS), zip -r $(DEP).ez $(DEP);) 35 | 36 | compile: $(DEPS_DIR) 37 | ./rebar compile -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # RabbitMQ Riak Exchange 2 | 3 | Latest tagged version works with RabbitMQ 3.3.5 and Riak 2.0. 4 | 5 | This is a custom exchange type for RabbitMQ that will put any message sent to it into Riak. 6 | By default, the Riak exchange will use your exchange name as the bucket name and your routing key as the Riak 7 | key. To alter this behaviour, simply pass special message headers: 8 | 9 | * `X-Riak-Bucket` - set this to the bucket name to use. 10 | * `X-Riak-Key` - set this to the key to use. 11 | 12 | For example, setting AMQP messages headers like the following: 13 | 14 | { 15 | "X-Riak-Bucket": "riak-test", 16 | "X-Riak-Key": "mykey" 17 | } 18 | 19 | Would result in a Riak object being stored at `/riak-test/mykey` 20 | 21 | ## Installation 22 | 23 | To install from source: 24 | 25 | git clone https://github.com/jbrisbin/riak-exchange 26 | cd riak-exchange 27 | make deps 28 | make 29 | make package 30 | cp deps/*.ez $RABBITMQ_HOME/plugins 31 | cp dist/*.ez $RABBITMQ_HOME/plugins 32 | 33 | Starting with RabbitMQ version 2.7, you must also enable the plugin using the new `rabbitmq-plugins` script. 34 | 35 | Issuing... 36 | 37 | rabbitmq-plugins list 38 | 39 | ...should give you a list of the plugins available to enable: 40 | 41 | [ ] protobuffs 0.8.1p4 42 | [ ] rabbit_exchange_type_riak 0.2.0 43 | [ ] riak_pb 2.0.0.16 44 | [ ] riakc 1.4.1-200-gb96d050 45 | 46 | Enable the plugin (including dependencies) by executing: 47 | 48 | rabbitmq-plugins enable rabbit_exchange_type_riak 49 | 50 | If you run `list` again, you should see them enabled: 51 | 52 | [e] protobuffs 0.8.1p4 53 | [E] rabbit_exchange_type_riak 0.2.0 54 | [e] riak_pb 2.0.0.16 55 | [e] riakc 1.4.1-200-gb96d050 56 | 57 | *NOTE:* I've also put up a tar file of the required .ez files you need to install in your RabbitMQ's plugins directory. 58 | 59 | ## Configuration 60 | 61 | To use the Riak exchange type, declare your exchange as type "x-riak". In addition to forwarding messages to 62 | Riak, this also acts like a regular exchange so you can have consumers bound to this exchange and they will 63 | receive the messages as well as going to Riak. 64 | 65 | To configure what Riak server to connect to, pass some arguments to the exchange declaration: 66 | 67 | * `host` - Hostname or IP of the Riak server to connect to. 68 | * `port` - Port number of the Riak server to connect to. 69 | * `maxclients` - The maximum number of clients to create in the pool (use more clients for higher-traffic exchanges). 70 | 71 | NEW in version 0.1.5: The Riak exchange can act like any valid RabbitMQ exchange type. Set an argument on your 72 | exchange when you declare it named `type_module` and give it a valid RabbitMQ exchange type module. You can use 73 | the ones built in that come with RabbitMQ, or you can use any custom ones you or a third party have written. 74 | 75 | Here's a sample list of possible exchange types and what you should set the `type_module` value to: 76 | 77 | * `direct`: `rabbit_exchange_type_direct` 78 | * `fanout`: `rabbit_exchange_type_fanout` 79 | * `headers`: `rabbit_exchange_type_headers` 80 | * `topic`: `rabbit_exchange_type_topic` 81 | * `random`: [`rabbit_exchange_type_random`](https://github.com/jbrisbin/random-exchange) 82 | 83 | If you don't specify anything, the exchange will default to a topic exchange. 84 | 85 | ## Metadata 86 | 87 | #### Content-Type 88 | 89 | Whatever Content-Type you set on the message will be the Content-Type used for Riak. 90 | 91 | #### Headers 92 | 93 | The Riak exchange will also translate your custom AMQP message headers into the special `X-Riak-Meta-` 94 | headers required by Riak. For example, if you set a custom header with the name `customheader`, your Riak 95 | document will have a header named `X-Riak-Meta-customheader`. 96 | -------------------------------------------------------------------------------- /src/rabbit_exchange_type_riak.erl: -------------------------------------------------------------------------------- 1 | -module(rabbit_exchange_type_riak). 2 | -include("riak_exchange.hrl"). 3 | -behaviour(rabbit_exchange_type). 4 | 5 | -define(BUCKET, <<"X-Riak-Bucket">>). 6 | -define(EXCHANGE_TYPE_BIN, <<"x-riak">>). 7 | -define(HOST, <<"host">>). 8 | -define(KEY, <<"X-Riak-Key">>). 9 | -define(PORT, <<"port">>). 10 | -define(MAX_CLIENTS, <<"maxclients">>). 11 | -define(TYPE, <<"type-module">>). 12 | 13 | -rabbit_boot_step({?MODULE, [ 14 | {description, "exchange type riak"}, 15 | {mfa, {rabbit_registry, register, [exchange, ?EXCHANGE_TYPE_BIN, ?MODULE]}}, 16 | {requires, rabbit_registry}, 17 | {enables, kernel_ready} 18 | ]}). 19 | 20 | -export([ 21 | add_binding/3, 22 | assert_args_equivalence/2, 23 | create/2, 24 | delete/3, 25 | policy_changed/2, 26 | description/0, 27 | recover/2, 28 | remove_bindings/3, 29 | route/2, 30 | serialise_events/0, 31 | validate/1, 32 | validate_binding/2 33 | ]). 34 | 35 | description() -> 36 | [{name, ?EXCHANGE_TYPE_BIN}, {description, <<"exchange type riak">>}]. 37 | 38 | serialise_events() -> 39 | false. 40 | 41 | validate(X) -> 42 | % io:format("Validate passed: ~w~n", [X]), 43 | Exchange = exchange_type(X), 44 | Exchange:validate(X). 45 | validate_binding(X, B) -> 46 | Exchange = exchange_type(X), 47 | Exchange:validate_binding(X, B). 48 | 49 | create(Tx, X = #exchange{name = #resource{virtual_host=_VirtualHost, name=_Name}, arguments = _Args}) -> 50 | XA = exchange_a(X), 51 | pg2:create(XA), 52 | 53 | case get_riak_client(X) of 54 | {ok, _Client} -> 55 | Exchange = exchange_type(X), 56 | Exchange:create(Tx, X); 57 | _ -> 58 | error_logger:error_msg("Could not connect to Riak"), 59 | {error, "could not connect to riak"} 60 | end. 61 | 62 | recover(X, _Bs) -> 63 | create(none, X). 64 | 65 | delete(Tx, X, Bs) -> 66 | XA = exchange_a(X), 67 | pg2:delete(XA), 68 | Exchange = exchange_type(X), 69 | Exchange:delete(Tx, X, Bs). 70 | 71 | policy_changed(X1, X2) -> 72 | Exchange = exchange_type(X1), 73 | Exchange:policy_changed(X1, X2). 74 | 75 | add_binding(Tx, X, B) -> 76 | Exchange = exchange_type(X), 77 | Exchange:add_binding(Tx, X, B). 78 | 79 | remove_bindings(Tx, X, Bs) -> 80 | Exchange = exchange_type(X), 81 | Exchange:remove_bindings(Tx, X, Bs). 82 | 83 | assert_args_equivalence(X, Args) -> 84 | rabbit_exchange:assert_args_equivalence(X, Args). 85 | 86 | route(X=#exchange{name = #resource{virtual_host = _VirtualHost, name = Name}}, 87 | D=#delivery{message = _Message0 = #basic_message{routing_keys = Routes, content = Content0}}) -> 88 | #content{ 89 | properties = _Props = #'P_basic'{ 90 | content_type = CT, 91 | headers = Headers, 92 | reply_to = _ReplyTo 93 | }, 94 | payload_fragments_rev = PayloadRev 95 | } = rabbit_binary_parser:ensure_content_decoded(Content0), 96 | 97 | ContentType = case CT of 98 | undefined -> <<"application/octet-stream">>; 99 | _ -> CT 100 | end, 101 | 102 | case get_riak_client(X) of 103 | {ok, Client} -> 104 | % Convert payload to list, concat together 105 | Payload = lists:foldl(fun(Chunk, NewPayload) -> 106 | <> 107 | end, <<>>, PayloadRev), 108 | % io:format("payload: ~p~n", [Payload]), 109 | % io:format("routes: ~p~n", [Routes]), 110 | 111 | lists:foldl(fun(Route, _) -> 112 | % Look for bucket from headers or default to exchange name 113 | Bucket = case Headers of 114 | undefined -> Name; 115 | _ -> case lists:keyfind(?BUCKET, 1, Headers) of 116 | {?BUCKET, _, B} -> B; 117 | _ -> Name 118 | end 119 | end, 120 | % Look for key from headers or default to routing key 121 | Key = case Headers of 122 | undefined -> Route; 123 | _ -> case lists:keyfind(?KEY, 1, Headers) of 124 | {?KEY, _, K} -> K; 125 | _ -> Route 126 | end 127 | end, 128 | 129 | % Insert or update everything 130 | % io:format("storing message: /~s/~s as ~s~n", [Bucket, Key, ContentType]), 131 | Obj0 = case riakc_pb_socket:get(Client, Bucket, Key) of 132 | {ok, OldObj} -> riakc_obj:update_value(OldObj, Payload, binary_to_list(ContentType)); 133 | _ -> riakc_obj:new(Bucket, Key, Payload, binary_to_list(ContentType)) 134 | end, 135 | 136 | % Populate metadata from msg properties 137 | Obj1 = case Headers of 138 | undefined -> Obj0; 139 | _ -> case lists:foldl(fun({PropKey, _PropType, PropVal}, NewProps) -> 140 | % io:format("key, type, val= (~p, ~p, ~p)~n", [PropKey, PropType, PropVal]), 141 | case PropKey of 142 | <<"X-Riak-Bucket", _/binary>> -> NewProps; 143 | <<"X-Riak-Key", _/binary>> -> NewProps; 144 | _ -> [{<<"X-Riak-Meta-", PropKey/binary>>, PropVal} | NewProps] 145 | end 146 | end, [], Headers) of 147 | [] -> Obj0; 148 | CMeta -> riakc_obj:update_metadata(Obj0, dict:store(<<"X-Riak-Meta">>, CMeta, riakc_obj:get_update_metadata(Obj0))) 149 | end 150 | end, 151 | 152 | % Insert/Update data 153 | _Result = riakc_pb_socket:put(Client, Obj1) 154 | % io:format("result: ~p~n", [Result]) 155 | end, [], Routes); 156 | _Err -> 157 | %io:format("err: ~p~n", [Err]), 158 | error_logger:error_msg("Could not connect to Riak") 159 | end, 160 | Exchange = exchange_type(X), 161 | Exchange:route(X, D). 162 | 163 | exchange_a(#exchange{name = #resource{virtual_host=VirtualHost, name=Name}}) -> 164 | list_to_atom(lists:flatten(io_lib:format("~s ~s", [VirtualHost, Name]))). 165 | 166 | get_riak_client(X=#exchange{arguments = Args}) -> 167 | Host = case lists:keyfind(?HOST, 1, Args) of 168 | {_, _, H} -> binary_to_list(H); 169 | _ -> "127.0.0.1" 170 | end, 171 | Port = case lists:keyfind(?PORT, 1, Args) of 172 | {_, _, P} -> 173 | {Pn, _} = string:to_integer(binary_to_list(P)), 174 | Pn; 175 | _ -> 8087 176 | end, 177 | MaxClients = case lists:keyfind(?MAX_CLIENTS, 1, Args) of 178 | {_, _, MC} -> 179 | {MCn, _} = string:to_integer(binary_to_list(MC)), 180 | MCn; 181 | _ -> 5 182 | end, 183 | XA = exchange_a(X), 184 | 185 | try 186 | case pg2:get_closest_pid(XA) of 187 | {error, _} -> create_riak_client(XA, Host, Port, MaxClients); 188 | PbClient -> 189 | case riakc_pb_socket:ping(PbClient) of 190 | pong -> {ok, PbClient}; 191 | _ -> 192 | error_logger:error_report("Disconnected Riak client discarded."), 193 | pg2:leave(XA, PbClient), 194 | get_riak_client(X) 195 | end 196 | end 197 | catch 198 | _ -> create_riak_client(XA, Host, Port, MaxClients) 199 | end. 200 | 201 | create_riak_client(XA, Host, Port, MaxClients) -> 202 | error_logger:info_report(io_lib:format("Starting ~p Riak PB clients to ~p:~p", [MaxClients, Host, Port])), 203 | case riakc_pb_socket:start_link(Host, Port) of 204 | {ok, PbClient} -> 205 | pg2:join(XA, PbClient), 206 | case length(pg2:get_members(XA)) of 207 | S when (S < MaxClients) -> create_riak_client(XA, Host, Port, MaxClients); 208 | _ -> {ok, PbClient} 209 | end; 210 | Err -> Err 211 | end. 212 | 213 | exchange_type(_Exchange=#exchange{ arguments=Args }) -> 214 | % io:format("Ensuring exchange type doesn't loop: ~w~n", [Exchange]), 215 | case lists:keyfind(?TYPE, 1, Args) of 216 | {?TYPE, _, Type} -> 217 | io:format("found type ~p~n", [Type]), 218 | case list_to_atom(binary_to_list(Type)) of 219 | rabbit_exchange_type_riak -> 220 | error_logger:error_report("Cannot base a Riak exchange on a Riak exchange. An infinite loop would occur."), 221 | rabbit_exchange_type_topic; 222 | Else -> Else 223 | end; 224 | _ -> rabbit_exchange_type_topic 225 | end. 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | --------------------------------------------------------------------------------