Implement a system for caching prepared SQL statements.

Prepared statements created via this mechanism are placed into the
cache for reuse once the user is done with them.  Statements won't
be reused as long as at least one proxy object still exists.  This
gives us a mechanism for safe reuse of statements without having to
worry that two different pieces of code might use the same statement
object and step on each other.

Also, boost::multi_index is awesome.

3 files changed, 238 insertions, 0 deletions

diff --git a/src/generic/util/sqlite.cc b/src/generic/util/sqlite.cc
index 4aaa3c1e..84a09b25 100644
--- a/src/generic/util/sqlite.cc
+++ b/src/generic/util/sqlite.cc
@@ -24,6 +24,11 @@ namespace aptitude
 {
   namespace sqlite
   {
+    /** \brief The default maximum size of a database's statement
+     *  cache.
+     */
+    const unsigned int default_statement_cache_limit = 100;
+
     db::lock::lock(db &parent)
       : handle(parent.handle)
     {
@@ -38,6 +43,7 @@ namespace aptitude
     db::db(const std::string &filename,
 	   int flags,
 	   const char *vfs)
+      : statement_cache_limit(default_statement_cache_limit)
     {
       const int result =
 	sqlite3_open_v2(filename.c_str(), &handle,
@@ -96,6 +102,59 @@ namespace aptitude
       return rval;
     }
 
+    void db::cache_statement(const statement_cache_entry &entry)
+    {
+      statement_cache_mru &mru(get_cache_mru());
+      mru.push_back(entry);
+
+      // Drop old entries from the cache if it's too large.
+      while(mru.size() > statement_cache_limit)
+	mru.pop_front();
+    }
+
+    db::statement_proxy_impl::~statement_proxy_impl()
+    {
+      // Careful here: the database might have been deleted while the
+      // proxy is active.  WE RELY ON THE FACT THAT DELETING THE
+      // DATABASE NULLS OUT THE STATEMENT HANDLE.
+      if(entry.stmt->handle == NULL)
+	return; // The database is dead; nothing to do.
+      else
+	entry.stmt->parent.cache_statement(entry);
+    }
+
+    db::statement_proxy db::get_cached_statement(const std::string &sql)
+    {
+      // Check whether the statement exists in the cache.
+      statement_cache_hash_index &index(get_cache_hash_index());
+
+      statement_cache_hash_index::const_iterator found =
+	index.find(sql);
+
+      if(found != index.end())
+	{
+	  // Extract the element from the set and return it.
+	  statement_cache_entry entry(*found);
+	  entry.stmt->reset();
+
+	  index.erase(sql);
+
+	  boost::shared_ptr<statement_proxy_impl> rval(new statement_proxy_impl(entry));
+	  return statement_proxy(rval);
+	}
+      else
+	{
+	  // Prepare a new SQL statement and return a proxy to it.  It
+	  // won't be added to the cache until the caller is done with
+	  // it.
+	  boost::shared_ptr<statement> stmt(statement::prepare(*this, sql));
+
+	  statement_cache_entry entry(sql, stmt);
+	  boost::shared_ptr<statement_proxy_impl> rval(new statement_proxy_impl(entry));
+	  return statement_proxy(rval);
+	}
+    }
+
 
 
     statement::statement(db &_parent, sqlite3_stmt *_handle)
diff --git a/src/generic/util/sqlite.h b/src/generic/util/sqlite.h
index ca2b9572..4aa11891 100644
--- a/src/generic/util/sqlite.h
+++ b/src/generic/util/sqlite.h
@@ -24,6 +24,10 @@
 
 #include <sqlite3.h>
 
+#include <boost/multi_index_container.hpp>
+#include <boost/multi_index/hashed_index.hpp>
+#include <boost/multi_index/sequenced_index.hpp>
+#include <boost/multi_index/member.hpp>
 #include <boost/shared_ptr.hpp>
 #include <boost/unordered_set.hpp>
 
@@ -76,6 +80,83 @@ namespace aptitude
       // Similarly, a set of active blob objects.
       boost::unordered_set<blob *> active_blobs;
 
+
+
+      // Used to cache statements for reuse.  Each statement in this
+      // set is currently *unused*; when a statement is requested, the
+      // requester effectively "checks out" a copy, removing it from
+      // the set.  The statement is accessed through a smart pointer
+      // wrapper that places it back in the set once it's no longer
+      // used.  This is necessary because it's not safe to reuse
+      // SQLite statements.
+      struct statement_cache_entry
+      {
+	std::string sql;
+	boost::shared_ptr<statement> stmt;
+
+	statement_cache_entry(const std::string &_sql,
+			      const boost::shared_ptr<statement> &_stmt)
+	  : sql(_sql), stmt(_stmt)
+	{
+	}
+      };
+
+      typedef boost::multi_index_container<
+	statement_cache_entry,
+	boost::multi_index::indexed_by<
+	  boost::multi_index::hashed_non_unique<
+	    boost::multi_index::member<
+	      statement_cache_entry,
+	      std::string,
+	      &statement_cache_entry::sql> >,
+	  boost::multi_index::sequenced<>
+	  >
+	>  statement_cache_container;
+
+      static const int statement_cache_hash_index_N = 0;
+      static const int statement_cache_mru_N = 1;
+
+      typedef statement_cache_container::nth_index<statement_cache_hash_index_N>::type statement_cache_hash_index;
+      typedef statement_cache_container::nth_index<statement_cache_mru_N>::type statement_cache_mru;
+
+      statement_cache_container statement_cache;
+      unsigned int statement_cache_limit;
+
+      statement_cache_hash_index &get_cache_hash_index()
+      {
+	return statement_cache.get<statement_cache_hash_index_N>();
+      }
+
+      statement_cache_mru &get_cache_mru()
+      {
+	return statement_cache.get<statement_cache_mru_N>();
+      }
+
+      void cache_statement(const statement_cache_entry &entry);
+
+
+      /** \brief An intermediate data item used to track the use of
+       *  a statement that was checked out from the cache.
+       *
+       *  When a statement_proxy is destroyed, it places its enclosed
+       *  statement back into the database's statement cache.
+       */
+      class statement_proxy_impl
+      {
+	statement_cache_entry entry;
+
+      public:
+	statement_proxy_impl(const statement_cache_entry &_entry)
+	  : entry(_entry)
+	{
+	}
+
+	const boost::shared_ptr<statement> &get_statement() const { return entry.stmt; }
+	const statement_cache_entry &get_entry() const { return entry; }
+
+	~statement_proxy_impl();
+      };
+
       db(const std::string &filename, int flags, const char *vfs);
     public:
       /** \brief Used to make the wrapper routines atomic.
@@ -114,6 +195,15 @@ namespace aptitude
       /** \brief Close the encapsulated database. */
       ~db();
 
+      /** \brief Change the maximum number of statements to cache.
+       *
+       *  If it is not set, the maximum number defaults to 100.
+       */
+      void set_statement_cache_limit(unsigned int new_limit)
+      {
+	statement_cache_limit = new_limit;
+      }
+
       /** \brief Retrieve the last error that was generated on this
        *  database.
        *
@@ -122,6 +212,49 @@ namespace aptitude
        *  another thread.
        */
       std::string get_error();
+
+      /** \brief Represents a statement retrieved from the
+       *  cache.
+       *
+       *  statement_proxy objects act as strong references to the
+       *  particular statement that was retrieved.  When all the
+       *  references to a statement expire, it is returned to the
+       *  cache as the most recently used entry.
+       *
+       *  Because statements are not thread-safe, statement proxies
+       *  should not be passed between threads (more specifically,
+       *  they should not be dereferenced from multiple threads at
+       *  once).  Instead, each thread should invoke
+       *  get_cached_statement() separately.
+       */
+      class statement_proxy
+      {
+	boost::shared_ptr<statement_proxy_impl> impl;
+
+	friend class db;
+
+	statement_proxy(const boost::shared_ptr<statement_proxy_impl> &_impl)
+	  : impl(_impl)
+	{
+	}
+
+      public:
+	statement &operator*() const { return *impl->get_statement(); }
+	statement *operator->() const { return impl->get_statement().get(); }
+
+	/** \brief Discard the reference to the implementation. */
+	void reset() { impl.reset(); }
+	/** \brief Test whether we have a valid pointer to the implementation. */
+	bool valid() const { return impl.get() != NULL; }
+      };
+
+      /** \brief Retrieve a statement from this database's statement
+       *  cache.
+       *
+       *  If the statement is not in the cache, it will be compiled
+       *  and added.
+       */
+      statement_proxy get_cached_statement(const std::string &sql);
     };
 
     /** \brief Wraps a prepared sqlite3 statement.
@@ -144,6 +277,7 @@ namespace aptitude
       bool has_data;
 
       friend class db;
+      friend class db::statement_proxy_impl;
 
       statement(db &_parent, sqlite3_stmt *_handle);
 
diff --git a/tests/test_sqlite.cc b/tests/test_sqlite.cc
index 598bd194..ce2f62b6 100644
--- a/tests/test_sqlite.cc
+++ b/tests/test_sqlite.cc
@@ -167,3 +167,48 @@ BOOST_FIXTURE_TEST_CASE(testGetString, test_db_fixture)
   BOOST_CHECK_THROW(stmt->get_string(0),
 		    exception);
 }
+
+BOOST_FIXTURE_TEST_CASE(testGetCachedStatement, memory_db_fixture)
+{
+  tmpdb->set_statement_cache_limit(2);
+  db::statement_proxy p1(tmpdb->get_cached_statement("create table foo(bar int)"));
+  db::statement_proxy p2(tmpdb->get_cached_statement("create table foo(bar int)"));
+  p2.reset();
+  p1.reset();
+
+  db::statement_proxy p3(tmpdb->get_cached_statement("create table foo(bar int)"));
+  db::statement_proxy p4(tmpdb->get_cached_statement("create table bar(foo int)"));
+
+  // Test that statements are being reused.
+  statement * const stmt1(&*p4);
+
+  p3.reset();
+  p4.reset();
+
+  db::statement_proxy p5(tmpdb->get_cached_statement("create table bar(foo int)"));
+
+  statement * const stmt2(&*p5);
+
+  BOOST_CHECK_EQUAL(stmt1, stmt2);
+
+  // Use the statements for some trivial operations.
+  db::statement_proxy p6(tmpdb->get_cached_statement("create table foo(bar int)"));
+  p6->exec();
+
+  db::statement_proxy p7(tmpdb->get_cached_statement("insert into foo (bar) values (5)"));
+  db::statement_proxy p8(tmpdb->get_cached_statement("select bar from foo"));
+
+  p7->exec();
+
+  BOOST_REQUIRE(p8->step());
+  BOOST_CHECK_EQUAL(p8->get_int(0), 5);
+  BOOST_CHECK(!p8->step());
+}
+
+BOOST_FIXTURE_TEST_CASE(getCachedStatementFail, memory_db_fixture)
+{
+  BOOST_REQUIRE_THROW(statement::prepare(*tmpdb, "select * from bar"),
+		      exception);
+  db::statement_proxy p1(tmpdb->get_cached_statement("create table foo(bar int)"));
+  db::statement_proxy p2(tmpdb->get_cached_statement("create table foo(bar int)"));
+}