From: Holger Weiss <holger@zedat.fu-berlin.de>
Date: Mon, 23 May 2016 22:25:52 +0000 (+0200)
Subject: mod_mam_mnesia: Don't exceed table size limit
X-Git-Tag: 16.06~85
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=2a9dd548b53e459afdb02a5c905199a152e4c01a;p=ejabberd

mod_mam_mnesia: Don't exceed table size limit

Don't write MAM messages into an Mnesia archive if the size of the table
comes close to the 2 GB limit for tables with disc-only copies.  That
way, the table is at least not corrupted when the limit is reached.
---

diff --git a/src/mod_mam_mnesia.erl b/src/mod_mam_mnesia.erl
index 1e1f20c9a..225d1c4ce 100644
--- a/src/mod_mam_mnesia.erl
+++ b/src/mod_mam_mnesia.erl
@@ -26,6 +26,8 @@
 	((A < B andalso byte_size(A) == byte_size(B))
 	 orelse byte_size(A) < byte_size(B))).
 
+-define(TABLE_SIZE_LIMIT, 2000000000). % A bit less than 2 GiB.
+
 %%%===================================================================
 %%% API
 %%%===================================================================
@@ -83,26 +85,35 @@ extended_fields() ->
     [].
 
 store(Pkt, _, {LUser, LServer}, Type, Peer, Nick, _Dir) ->
-    LPeer = {PUser, PServer, _} = jid:tolower(Peer),
-    TS = p1_time_compat:timestamp(),
-    ID = jlib:integer_to_binary(now_to_usec(TS)),
-    F = fun() ->
-		mnesia:write(#archive_msg{us = {LUser, LServer},
-					  id = ID,
-					  timestamp = TS,
-					  peer = LPeer,
-					  bare_peer = {PUser, PServer, <<>>},
-					  type = Type,
-					  nick = Nick,
-					  packet = Pkt})
-	end,
-    case mnesia:transaction(F) of
-	{atomic, ok} ->
-	    {ok, ID};
-	{aborted, Err} ->
-	    ?ERROR_MSG("Cannot add message to MAM archive of ~s@~s: ~s",
-		       [LUser, LServer, Err]),
-	    Err
+    case {mnesia:table_info(archive_msg, disc_only_copies),
+	  mnesia:table_info(archive_msg, memory)} of
+	{[_|_], TableSize} when TableSize > ?TABLE_SIZE_LIMIT ->
+	    ?ERROR_MSG("MAM archive of ~s@~s too large, won't store message",
+		       [LUser, LServer]),
+	    {error, overflow};
+	_ ->
+	    LPeer = {PUser, PServer, _} = jid:tolower(Peer),
+	    TS = p1_time_compat:timestamp(),
+	    ID = jlib:integer_to_binary(now_to_usec(TS)),
+	    F = fun() ->
+			mnesia:write(
+			  #archive_msg{us = {LUser, LServer},
+				       id = ID,
+				       timestamp = TS,
+				       peer = LPeer,
+				       bare_peer = {PUser, PServer, <<>>},
+				       type = Type,
+				       nick = Nick,
+				       packet = Pkt})
+		end,
+	    case mnesia:transaction(F) of
+		{atomic, ok} ->
+		    {ok, ID};
+		{aborted, Err} ->
+		    ?ERROR_MSG("Cannot add message to MAM archive of ~s@~s: ~s",
+			       [LUser, LServer, Err]),
+		    Err
+	    end
     end.
 
 write_prefs(_LUser, _LServer, Prefs, _ServerHost) ->