From 2ae0bdafbb58c3e4f8ee190a5f9bd3f61a70acb1 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Sun, 7 Jun 2015 13:15:13 -0700
Subject: [PATCH] memdb: testing writer concurrency

---
 nomad/memdb/memdb.go      | 12 ++++++------
 nomad/memdb/memdb_test.go | 41 +++++++++++++++++++++++++++++++++++++++
 nomad/memdb/txn.go        | 11 +++++------
 3 files changed, 52 insertions(+), 12 deletions(-)
 create mode 100644 nomad/memdb/memdb_test.go

diff --git a/nomad/memdb/memdb.go b/nomad/memdb/memdb.go
index 89465983a..199732f2c 100644
--- a/nomad/memdb/memdb.go
+++ b/nomad/memdb/memdb.go
@@ -36,13 +36,13 @@ func NewMemDB(schema *DBSchema) (*MemDB, error) {
 // Txn is used to start a new transaction, in either read or write mode.
 // There can only be a single concurrent writer, but any number of readers.
 func (db *MemDB) Txn(write bool) *Txn {
-	txn := &Txn{
-		db:    db,
-		write: write,
-		root:  db.root,
-	}
 	if write {
-		txn.rootTxn = txn.root.Txn()
+		db.writer.Lock()
+	}
+	txn := &Txn{
+		db:      db,
+		write:   write,
+		rootTxn: db.root.Txn(),
 	}
 	return txn
 }
diff --git a/nomad/memdb/memdb_test.go b/nomad/memdb/memdb_test.go
new file mode 100644
index 000000000..134713d81
--- /dev/null
+++ b/nomad/memdb/memdb_test.go
@@ -0,0 +1,41 @@
+package memdb
+
+import (
+	"testing"
+	"time"
+)
+
+func TestMemDB_SingleWriter_MultiReader(t *testing.T) {
+	db, err := NewMemDB(testValidSchema())
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	tx1 := db.Txn(true)
+	tx2 := db.Txn(false) // Should not block!
+	tx3 := db.Txn(false) // Should not block!
+	tx4 := db.Txn(false) // Should not block!
+
+	doneCh := make(chan struct{})
+	go func() {
+		defer close(doneCh)
+		db.Txn(true)
+	}()
+
+	select {
+	case <-doneCh:
+		t.Fatalf("should not allow another writer")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	tx1.Abort()
+	tx2.Abort()
+	tx3.Abort()
+	tx4.Abort()
+
+	select {
+	case <-doneCh:
+	case <-time.After(10 * time.Millisecond):
+		t.Fatalf("should allow another writer")
+	}
+}
diff --git a/nomad/memdb/txn.go b/nomad/memdb/txn.go
index 42091f705..371be50fa 100644
--- a/nomad/memdb/txn.go
+++ b/nomad/memdb/txn.go
@@ -10,7 +10,6 @@ import (
 type Txn struct {
 	db      *MemDB
 	write   bool
-	root    *iradix.Tree
 	rootTxn *iradix.Txn
 }
 
@@ -22,14 +21,15 @@ func (txn *Txn) Abort() {
 	}
 
 	// Check if already aborted or committed
-	if txn.root == nil {
+	if txn.rootTxn == nil {
 		return
 	}
 
+	// Clear the txn
+	txn.rootTxn = nil
+
 	// Release the writer lock since this is invalid
 	txn.db.writer.Unlock()
-	txn.root = nil
-	txn.rootTxn = nil
 }
 
 // Commit is used to finalize this transaction. This is a noop for read transactions.
@@ -40,7 +40,7 @@ func (txn *Txn) Commit() {
 	}
 
 	// Check if already aborted or committed
-	if txn.root == nil {
+	if txn.rootTxn == nil {
 		return
 	}
 
@@ -48,7 +48,6 @@ func (txn *Txn) Commit() {
 	txn.db.root = txn.rootTxn.Commit()
 
 	// Clear the txn
-	txn.root = nil
 	txn.rootTxn = nil
 
 	// Release the writer lock since this is invalid