From f34a9350c3c7c815ccf9d8d1dde10227a03ce8bc Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Tue, 22 Mar 2016 11:59:28 -0700
Subject: [PATCH] benchmark: benchmark stm workload with distributed mutex

---
 tools/benchmark/cmd/stm.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tools/benchmark/cmd/stm.go b/tools/benchmark/cmd/stm.go
index 1b1d97ad3..407df9246 100644
--- a/tools/benchmark/cmd/stm.go
+++ b/tools/benchmark/cmd/stm.go
@@ -45,6 +45,7 @@ var (
 	stmKeyCount     int
 	stmValSize      int
 	stmWritePercent int
+	stmMutex        bool
 	mkSTM           func(context.Context, *v3.Client, func(v3sync.STM) error) (*v3.TxnResponse, error)
 )
 
@@ -56,6 +57,7 @@ func init() {
 	stmCmd.Flags().IntVar(&stmTotal, "total", 10000, "Total number of completed STM transactions")
 	stmCmd.Flags().IntVar(&stmKeysPerTxn, "keys-per-txn", 1, "Number of keys to access per transaction")
 	stmCmd.Flags().IntVar(&stmWritePercent, "txn-wr-percent", 50, "Percentage of keys to overwrite per transaction")
+	stmCmd.Flags().BoolVar(&stmMutex, "use-mutex", false, "Wrap STM transaction in a distributed mutex")
 	stmCmd.Flags().IntVar(&stmValSize, "val-size", 8, "Value size of each STM put request")
 }
 
@@ -139,9 +141,20 @@ func stmFunc(cmd *cobra.Command, args []string) {
 func doSTM(ctx context.Context, client *v3.Client, requests <-chan stmApply) {
 	defer wg.Done()
 
+	var m *v3sync.Mutex
+	if stmMutex {
+		m = v3sync.NewMutex(client, "stmlock")
+	}
+
 	for applyf := range requests {
 		st := time.Now()
+		if m != nil {
+			m.Lock(context.TODO())
+		}
 		_, err := mkSTM(context.TODO(), client, applyf)
+		if m != nil {
+			m.Unlock(context.TODO())
+		}
 		var errStr string
 		if err != nil {
 			errStr = err.Error()