From 7e1e95fa2a4433864788817c54743a25c8a363a2 Mon Sep 17 00:00:00 2001
From: rnhmjoj <rnhmjoj@inventati.org>
Date: Wed, 6 Nov 2019 21:33:03 +0100
Subject: [PATCH] save url table on exit

---
 src/Breve/UrlTable.hs | 19 +++++++++++++------
 src/Main.hs           | 10 ++++++++--
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/src/Breve/UrlTable.hs b/src/Breve/UrlTable.hs
index d97fbe0..aafb64a 100644
--- a/src/Breve/UrlTable.hs
+++ b/src/Breve/UrlTable.hs
@@ -5,6 +5,7 @@ to store the URLs in memory and on disk.
 module Breve.UrlTable
 ( UrlTable
 , load
+, save
 , insert
 , extract
 ) where
@@ -19,20 +20,26 @@ import qualified Data.HashTable.IO as H
 -- | The hash table that stores URLs
 type UrlTable = H.CuckooHashTable Name Url
 
--- | Periodically writes a 'UrlTable' to a file
---
--- The table is stored in a text file
--- as Haskell code for semplicity.
+-- | Periodically save a 'UrlTable' to a file
 sync :: UrlTable -> FilePath -> IO ()
 sync table file = forever $ do
   threadDelay (round 3.0e8)
+  save table file
+
+-- | Writes a 'UrlTable' to a file
+--
+-- The table is stored in a text file
+-- as Haskell code for semplicity.
+save :: UrlTable -> FilePath -> IO ()
+save table file = do
   content <- show <$> H.toList table
   writeFile file content
+  putStrLn "\n[breve] url table synced."
 
 -- | Loads a URL table from a file
 --
--- The format should be the same one used
--- by the 'sync' function.
+-- Once the file is loaded it will be synced
+-- periodically (every 5min) on the disk.
 load :: FilePath -> IO UrlTable
 load file = do
   content <- readFile file
diff --git a/src/Main.hs b/src/Main.hs
index d71418e..371b2c6 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -20,6 +20,7 @@ import Data.Maybe      (listToMaybe)
 
 -- IO
 import Control.Monad      (when, void)
+import Control.Exception  as E 
 import Control.Concurrent (forkIO)
 import System.Environment (getArgs)
 import Data.Text.IO       as T
@@ -61,9 +62,14 @@ main = do
   when (bindPort == 443) $ void $
     forkIO (run 80 $ forceSSL emptyApp)
 
+  -- Save the table just before exiting
+  let exit E.UserInterrupt = save table urlTable
+      exit e               = throwIO e
+
   -- Middlewares are functions (Application -> Application).
   -- We use one here to add requests
   let middlewares = logStdout
 
-  T.putStrLn ("Serving on " <> bindUrl)
-  runApp config (middlewares $ breve static bindUrl table)
+  handle exit $ do
+    T.putStrLn ("Serving on " <> bindUrl)
+    runApp config (middlewares $ breve static bindUrl table)