From 2f24bcdf4aa430d5514a5a051c9cd2aef6b44e07 Mon Sep 17 00:00:00 2001
From: Anthony Yeh <enisoc@google.com>
Date: Fri, 5 Jun 2015 02:22:42 -0700
Subject: [PATCH] Add compatibility flag to lock for all serving graph updates.

---
 go/vt/topotools/rebuild.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/go/vt/topotools/rebuild.go b/go/vt/topotools/rebuild.go
index 6be99e5506..4e0e76114d 100644
--- a/go/vt/topotools/rebuild.go
+++ b/go/vt/topotools/rebuild.go
@@ -5,6 +5,7 @@
 package topotools
 
 import (
+	"flag"
 	"fmt"
 	"sync"
 	"time"
@@ -17,6 +18,8 @@ import (
 	"golang.org/x/net/context"
 )
 
+var lockSrvShard = flag.Bool("lock_srvshard", true, "serialize serving graph updates by locking the SrvShard")
+
 // RebuildShard updates the SrvShard objects and underlying serving graph.
 //
 // Re-read from TopologyServer to make sure we are using the side
@@ -339,7 +342,20 @@ func retryUpdateEndpoints(ctx context.Context, ts topo.Server, cell, keyspace, s
 
 // UpdateTabletEndpoints fixes up any entries in the serving graph that relate
 // to a given tablet.
-func UpdateTabletEndpoints(ctx context.Context, ts topo.Server, tablet *topo.Tablet) error {
+func UpdateTabletEndpoints(ctx context.Context, ts topo.Server, tablet *topo.Tablet) (err error) {
+	if *lockSrvShard {
+		// This lock is only necessary until all tablets are upgraded to lock-free.
+		actionNode := actionnode.RebuildSrvShard()
+		lockPath, err := actionNode.LockSrvShard(ctx, ts, tablet.Alias.Cell, tablet.Keyspace, tablet.Shard)
+		if err != nil {
+			return fmt.Errorf("can't lock shard for UpdateTabletEndpoints(%v): %v", tablet, err)
+		}
+
+		defer func() {
+			actionNode.UnlockSrvShard(ctx, ts, tablet.Alias.Cell, tablet.Keyspace, tablet.Shard, lockPath, err)
+		}()
+	}
+
 	srvTypes, err := ts.GetSrvTabletTypesPerShard(ctx, tablet.Alias.Cell, tablet.Keyspace, tablet.Shard)
 	if err != nil {
 		if err != topo.ErrNoNode {