minor update

uber-archive · Jul 31, 2017 · 9607068 · 9607068
1 parent 5bf4ea2
commit 9607068
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/services/replicator/replicator.go b/services/replicator/replicator.go
@@ -253,14 +253,15 @@ func (r *Replicator) OpenReplicationReadStreamHandler(w http.ResponseWriter, req
 	outConn := newOutConnection(extUUID, destDesc.GetPath(), outStream, r.logger, r.m3Client, metrics.OpenReplicationReadScope)
 	outConn.open()
 	r.addStoreHostConn(outConn)
+	defer r.removeStoreHostConn(outConn)
 
 	inConn := newInConnection(extUUID, destDesc.GetPath(), inStream, outConn.msgsCh, r.logger, r.m3Client, metrics.OpenReplicationReadScope, metrics.OpenReplicationReadPerDestScope)
 	inConn.open()
 
 	go r.manageInOutConn(inConn, outConn)
 	<-inConn.closeChannel
 	<-outConn.closeChannel
-	r.removeStoreHostConn(outConn)
+
 	return
 }
 
@@ -333,14 +334,15 @@ func (r *Replicator) OpenReplicationRemoteReadStreamHandler(w http.ResponseWrite
 	outConn := newOutConnection(extUUID, destDesc.GetPath(), outStream, r.logger, r.m3Client, metrics.OpenReplicationRemoteReadScope)
 	outConn.open()
 	r.addRemoteReplicatorConn(outConn)
+	defer r.removeRemoteReplicatorConn(outConn)
 
 	inConn := newInConnection(extUUID, destDesc.GetPath(), inStream, outConn.msgsCh, r.logger, r.m3Client, metrics.OpenReplicationRemoteReadScope, metrics.OpenReplicationRemoteReadPerDestScope)
 	inConn.open()
 
 	go r.manageInOutConn(inConn, outConn)
 	<-inConn.closeChannel
 	<-outConn.closeChannel
-	r.removeRemoteReplicatorConn(outConn)
+
 	return
 }