Optimize a bunch of stuff

3 files changed, 42 insertions, 13 deletions

diff --git a/funcs.go b/funcs.go
index 4967e83..47c6273 100644
--- a/funcs.go
+++ b/funcs.go
@@ -1,14 +1,26 @@
 package stm
 
+import (
+	"sync"
+)
+
+var (
+	txPool = sync.Pool{New: func() interface{} {
+		tx := &Tx{
+			reads:  make(map[*Var]uint64),
+			writes: make(map[*Var]interface{}),
+		}
+		tx.cond.L = &globalLock
+		return tx
+	}}
+)
+
 // Atomically executes the atomic function fn.
 func Atomically(fn func(*Tx)) interface{} {
-retry:
 	// run the transaction
-	tx := &Tx{
-		reads:  make(map[*Var]uint64),
-		writes: make(map[*Var]interface{}),
-	}
-	tx.cond.L = &globalLock
+	tx := txPool.Get().(*Tx)
+retry:
+	tx.reset()
 	var ret interface{}
 	if func() (retry bool) {
 		defer func() {
@@ -40,6 +52,7 @@ retry:
 	// commit the write log and broadcast that variables have changed
 	tx.commit()
 	globalLock.Unlock()
+	tx.recycle()
 	return ret
 }
 
diff --git a/tx.go b/tx.go
index 5dd41ae..bae2006 100644
--- a/tx.go
+++ b/tx.go
@@ -24,16 +24,17 @@ func (tx *Tx) verify() bool {
 	return true
 }
 
-// commit writes the values in the transaction log to their respective Vars.
+// Writes the values in the transaction log to their respective Vars.
 func (tx *Tx) commit() {
 	for v, val := range tx.writes {
 		v.mu.Lock()
 		v.val = val
 		v.version++
+		v.mu.Unlock()
 		for tx := range v.watchers {
 			tx.cond.Broadcast()
+			delete(v.watchers, tx)
 		}
-		v.mu.Unlock()
 	}
 }
 
@@ -95,3 +96,17 @@ func (tx *Tx) Return(v interface{}) {
 type _return struct {
 	value interface{}
 }
+
+func (tx *Tx) reset() {
+	for k := range tx.reads {
+		delete(tx.reads, k)
+	}
+	for k := range tx.writes {
+		delete(tx.writes, k)
+	}
+}
+
+func (tx *Tx) recycle() {
+	tx.reset()
+	txPool.Put(tx)
+}
diff --git a/var.go b/var.go
index 05e86cf..bf48a41 100644
--- a/var.go
+++ b/var.go
@@ -2,15 +2,16 @@ package stm
 
 import "sync"
 
-// A Var holds an STM variable.
+// Holds an STM variable.
 type Var struct {
-	val      interface{}
-	version  uint64
-	mu       sync.Mutex
+	mu      sync.Mutex
+	val     interface{}
+	version uint64
+
 	watchers map[*Tx]struct{}
 }
 
-// NewVar returns a new STM variable.
+// Returns a new STM variable.
 func NewVar(val interface{}) *Var {
 	return &Var{
 		val:      val,