diff --git a/clients/index.js b/clients/index.js index 9f105f4..f17e380 100644 --- a/clients/index.js +++ b/clients/index.js @@ -365,6 +365,7 @@ ApplicationClients.prototype.onRemoteConfigUpdate = function onRemoteConfigUpdat self.updateReservoir(); self.updateReapPeersPeriod(); self.updatePrunePeersPeriod(); + self.updateDrainIncomingConnections(); self.updatePartialAffinityEnabled(); self.setMaximumRelayTTL(); self.updatePeerHeapEnabled(); @@ -444,6 +445,12 @@ function updatePrunePeersPeriod() { self.serviceProxy.setPrunePeersPeriod(period); }; +ApplicationClients.prototype.updateDrainIncomingConnections = function updateDrainIncomingConnections() { + var self = this; + var enabled = self.remoteConfig.get('drainIncomingConnections.enabled', false); + self.serviceProxy.setDrainIncomingConnections(enabled); +}; + ApplicationClients.prototype.updatePartialAffinityEnabled = function updatePartialAffinityEnabled() { var self = this; var enabled = self.remoteConfig.get('partialAffinity.enabled', false); diff --git a/service-proxy.js b/service-proxy.js index 2a5a50e..10f3c8c 100644 --- a/service-proxy.js +++ b/service-proxy.js @@ -89,6 +89,7 @@ function ServiceDispatchHandler(options) { }); self.rateLimiterEnabled = options.rateLimiterEnabled; + self.drainIncomingConnections = !!options.drainIncomingConnections; self.partialAffinityEnabled = !!options.partialAffinityEnabled; self.minPeersPerWorker = options.minPeersPerWorker || DEFAULT_MIN_PEERS_PER_WORKER; self.minPeersPerRelay = options.minPeersPerRelay || DEFAULT_MIN_PEERS_PER_RELAY; @@ -635,7 +636,75 @@ function ensurePeerConnected(serviceName, peer, reason, now) { peer.clearDrain('canceled to ensure peer connection'); } - peer.connectTo(); + var conn = peer.connectTo(); + if (self.drainIncomingConnections) { + self.drainInOnceConnected(peer, conn, reason); + } +}; + +ServiceDispatchHandler.prototype.drainInOnceConnected = +function drainInOnceConnected(peer, conn, reason) { + var self = this; + + peer.waitForIdentified(conn, onConnIded); + + function onConnIded(err) { + if (err) { + self.logger.warn( + 'failed to ensure outgoing connection to service peer', + self.extendLogInfo({ + error: err, + peerHostPort: peer.hostPort, + refreshReason: reason + })); + return; + } + peer.drain({ + reason: reason, + direction: 'in', + timeout: self.drainTimeout + }, connectDrainDone); + } + + function connectDrainDone(err) { + if (err && + err.type === 'tchannel.drain.peer.timed-out') { + // TODO: stat? + self.logger.warn( + 'forcibly closing drained peer', + self.extendLogInfo({ + error: err, + drainReason: reason + }) + ); + err = null; + } + if (err) { + self.logger.warn( + 'failed to drain incoming connections from service peer', + self.extendLogInfo({ + error: err, + peerHostPort: peer.hostPort + }) + ); + peer.clearDrain(); + return; + } + peer.closeDrainedConnections(connectDrainCloseDone); + } + + function connectDrainCloseDone(err) { + if (err) { + self.logger.warn( + 'failed to close drained incoming connections from service peer', + self.extendLogInfo({ + error: err, + peerHostPort: peer.hostPort + }) + ); + } + peer.clearDrain(); + } }; ServiceDispatchHandler.prototype.getPartialRange = @@ -1538,6 +1607,12 @@ function disableRateLimiter() { self.rateLimiterEnabled = false; }; +ServiceDispatchHandler.prototype.setDrainIncomingConnections = +function setDrainIncomingConnections(enabled) { + var self = this; + self.drainIncomingConnections = !!enabled; +}; + ServiceDispatchHandler.prototype.setPartialAffinityEnabled = function enablePartialAffinity(enabled) { var self = this;