diff hgext3rd/topic/discovery.py @ 2676:10dedac0d82e

topic: also insert the extra head check with using the new head checking

This will get the protection in place in all cases.

--- 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