# HG changeset patch
# User Pierre-Yves David <pierre-yves.david@octobus.net>
# Date 1498663736 -7200
# Node ID 10dedac0d82e6cdd6c20a44a7a3e37875bc420c2
# Parent  304232cc14b631c1e64564ffba90f759343e098f
topic: also insert the extra head check with using the new head checking

This will get the protection in place in all cases.

diff -r 304232cc14b6 -r 10dedac0d82e hgext3rd/topic/discovery.py
--- a/hgext3rd/topic/discovery.py	Wed Jun 28 17:38:04 2017 +0200
+++ b/hgext3rd/topic/discovery.py	Wed Jun 28 17:28:56 2017 +0200
@@ -9,6 +9,7 @@
     error,
     exchange,
     extensions,
+    util,
     wireproto,
 )
 
@@ -147,9 +148,14 @@
     extensions.wrapfunction(discovery, '_headssummary', _headssummary)
     extensions.wrapfunction(wireproto, 'branchmap', wireprotobranchmap)
     extensions.wrapfunction(wireproto, '_capabilities', wireprotocaps)
+    # we need a proper wrap b2 part stuff
     extensions.wrapfunction(bundle2, 'handlecheckheads', handlecheckheads)
-    # we need a proper wrap b2 part stuff
     bundle2.handlecheckheads.params = frozenset()
     bundle2.parthandlermapping['check:heads'] = bundle2.handlecheckheads
+    if util.safehasattr(bundle2, 'handlecheckupdatedheads'):
+        # we still need a proper wrap b2 part stuff
+        extensions.wrapfunction(bundle2, 'handlecheckupdatedheads', handlecheckheads)
+        bundle2.handlecheckupdatedheads.params = frozenset()
+        bundle2.parthandlermapping['check:updated-heads'] = bundle2.handlecheckupdatedheads
     extensions.wrapfunction(exchange, '_pushb2phases', _pushb2phases)
     exchange.b2partsgenmapping['phase'] = exchange._pushb2phases