3 files changed, 167 insertions, 2 deletions

diff --git a/webkit/database/database_connections.cc b/webkit/database/database_connections.cc
index 05efb1e..612f73b4 100644
--- a/webkit/database/database_connections.cc
+++ b/webkit/database/database_connections.cc
@@ -1,10 +1,13 @@
-// Copyright (c) 2009 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "webkit/database/database_connections.h"
 
+#include "base/auto_reset.h"
 #include "base/logging.h"
+#include "base/message_loop.h"
+#include "base/message_loop_proxy.h"
 
 namespace webkit_database {
 
@@ -84,4 +87,53 @@ void DatabaseConnections::RemoveConnectionsHelper(
   }
 }
 
+DatabaseConnectionsWrapper::DatabaseConnectionsWrapper()
+    : waiting_for_dbs_to_close_(false),
+      main_thread_(base::MessageLoopProxy::CreateForCurrentThread()) {
+}
+
+DatabaseConnectionsWrapper::~DatabaseConnectionsWrapper() {
+}
+
+void DatabaseConnectionsWrapper::WaitForAllDatabasesToClose() {
+  // We assume that new databases won't be open while we're waiting.
+  DCHECK(main_thread_->BelongsToCurrentThread());
+  if (HasOpenConnections()) {
+    AutoReset<bool> auto_reset(&waiting_for_dbs_to_close_, true);
+    MessageLoop::ScopedNestableTaskAllower nestable(MessageLoop::current());
+    MessageLoop::current()->Run();
+  }
+}
+
+bool DatabaseConnectionsWrapper::HasOpenConnections() {
+  DCHECK(main_thread_->BelongsToCurrentThread());
+  base::AutoLock auto_lock(open_connections_lock_);
+  return !open_connections_.IsEmpty();
+}
+
+void DatabaseConnectionsWrapper::AddOpenConnection(
+    const string16& origin_identifier,
+    const string16& database_name) {
+  // We add to the collection immediately on any thread.
+  base::AutoLock auto_lock(open_connections_lock_);
+  open_connections_.AddConnection(origin_identifier, database_name);
+}
+
+void DatabaseConnectionsWrapper::RemoveOpenConnection(
+    const string16& origin_identifier,
+    const string16& database_name) {
+  // But only remove from the collection on the main thread
+  // so we can handle the waiting_for_dbs_to_close_ case.
+  if (!main_thread_->BelongsToCurrentThread()) {
+    main_thread_->PostTask(FROM_HERE, NewRunnableMethod(
+        this, &DatabaseConnectionsWrapper::RemoveOpenConnection,
+        origin_identifier, database_name));
+    return;
+  }
+  base::AutoLock auto_lock(open_connections_lock_);
+  open_connections_.RemoveConnection(origin_identifier, database_name);
+  if (waiting_for_dbs_to_close_ && open_connections_.IsEmpty())
+    MessageLoop::current()->Quit();
+}
+
 }  // namespace webkit_database
diff --git a/webkit/database/database_connections.h b/webkit/database/database_connections.h
index 433bb3f..6a594b2 100644
--- a/webkit/database/database_connections.h
+++ b/webkit/database/database_connections.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2009 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
@@ -8,7 +8,13 @@
 #include <map>
 #include <vector>
 
+#include "base/memory/ref_counted.h"
 #include "base/string16.h"
+#include "base/synchronization/lock.h"
+
+namespace base {
+class MessageLoopProxy;
+}
 
 namespace webkit_database {
 
@@ -40,6 +46,34 @@ class DatabaseConnections {
                                int num_connections);
 };
 
+// A wrapper class that provides thread-safety and the
+// ability to wait until all connections have closed.
+// Intended for use in renderer processes.
+class DatabaseConnectionsWrapper
+    : public base::RefCountedThreadSafe<DatabaseConnectionsWrapper> {
+ public:
+  DatabaseConnectionsWrapper();
+
+  // The Wait and Has methods should only be called on the
+  // main thread (the thread on which the wrapper is constructed).
+  void WaitForAllDatabasesToClose();
+  bool HasOpenConnections();
+
+  // Add and Remove may be called on any thread.
+  void AddOpenConnection(const string16& origin_identifier,
+                         const string16& database_name);
+  void RemoveOpenConnection(const string16& origin_identifier,
+                            const string16& database_name);
+ private:
+  ~DatabaseConnectionsWrapper();
+  friend class base::RefCountedThreadSafe<DatabaseConnectionsWrapper>;
+
+  bool waiting_for_dbs_to_close_;
+  base::Lock open_connections_lock_;
+  DatabaseConnections open_connections_;
+  scoped_refptr<base::MessageLoopProxy> main_thread_;
+};
+
 }  // namespace webkit_database
 
 #endif  // WEBKIT_DATABASE_DATABASE_CONNECTIONS_H_
diff --git a/webkit/database/database_connections_unittest.cc b/webkit/database/database_connections_unittest.cc
new file mode 100644
index 0000000..1b3ef34
--- /dev/null
+++ b/webkit/database/database_connections_unittest.cc
@@ -0,0 +1,79 @@
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "base/message_loop.h"
+#include "base/task.h"
+#include "base/threading/thread.h"
+#include "base/utf_string_conversions.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "webkit/database/database_connections.h"
+
+namespace webkit_database {
+
+namespace {
+
+void RemoveConnectionTask(
+    const string16& origin_id, const string16& database_name,
+    scoped_refptr<DatabaseConnectionsWrapper> obj,
+    bool* did_task_execute) {
+  *did_task_execute = true;
+  obj->RemoveOpenConnection(origin_id, database_name);
+}
+
+void ScheduleRemoveConnectionTask(
+    base::Thread* thread,  const string16& origin_id,
+    const string16& database_name,
+    scoped_refptr<DatabaseConnectionsWrapper> obj,
+    bool* did_task_execute) {
+  thread->message_loop()->PostTask(FROM_HERE,
+      NewRunnableFunction(RemoveConnectionTask,
+                          origin_id, database_name, obj, did_task_execute));
+}
+
+}  // anonymous namespace
+
+TEST(DatabaseConnectionsTest, DatabaseConnectionsWrapperTest) {
+  string16 kOriginId(ASCIIToUTF16("origin_id"));
+  string16 kName(ASCIIToUTF16("database_name"));
+
+  scoped_refptr<DatabaseConnectionsWrapper> obj(
+      new DatabaseConnectionsWrapper);
+  EXPECT_FALSE(obj->HasOpenConnections());
+  obj->AddOpenConnection(kOriginId, kName);
+  EXPECT_TRUE(obj->HasOpenConnections());
+  obj->AddOpenConnection(kOriginId, kName);
+  EXPECT_TRUE(obj->HasOpenConnections());
+  obj->RemoveOpenConnection(kOriginId, kName);
+  EXPECT_TRUE(obj->HasOpenConnections());
+  obj->RemoveOpenConnection(kOriginId, kName);
+  EXPECT_FALSE(obj->HasOpenConnections());
+  obj->WaitForAllDatabasesToClose();  // should return immediately
+
+  // Test WaitForAllDatabasesToClose with the last connection
+  // being removed on the current thread.
+  obj->AddOpenConnection(kOriginId, kName);
+  bool did_task_execute = false;
+  MessageLoop::current()->PostTask(FROM_HERE,
+      NewRunnableFunction(&RemoveConnectionTask,
+                          kOriginId, kName, obj, &did_task_execute));
+  obj->WaitForAllDatabasesToClose();  // should return after the task executes
+  EXPECT_TRUE(did_task_execute);
+  EXPECT_FALSE(obj->HasOpenConnections());
+
+  // Test WaitForAllDatabasesToClose with the last connection
+  // being removed on another thread.
+  obj->AddOpenConnection(kOriginId, kName);
+  base::Thread thread("WrapperTestThread");
+  thread.Start();
+  did_task_execute = false;
+  MessageLoop::current()->PostTask(FROM_HERE,
+      NewRunnableFunction(ScheduleRemoveConnectionTask,
+                          &thread, kOriginId, kName,
+                          obj, &did_task_execute));
+  obj->WaitForAllDatabasesToClose();  // should return after the task executes
+  EXPECT_TRUE(did_task_execute);
+  EXPECT_FALSE(obj->HasOpenConnections());
+}
+
+}  // namespace webkit_database