23 # common |
23 # common |
24 |
24 |
25 |
25 |
26 def _getmaster(ui): |
26 def _getmaster(ui): |
27 """get the mainbranch, and enforce it is set""" |
27 """get the mainbranch, and enforce it is set""" |
28 master = ui.config('fastannotate', 'mainbranch') |
28 master = ui.config(b'fastannotate', b'mainbranch') |
29 if not master: |
29 if not master: |
30 raise error.Abort( |
30 raise error.Abort( |
31 _( |
31 _( |
32 'fastannotate.mainbranch is required ' |
32 b'fastannotate.mainbranch is required ' |
33 'for both the client and the server' |
33 b'for both the client and the server' |
34 ) |
34 ) |
35 ) |
35 ) |
36 return master |
36 return master |
37 |
37 |
38 |
38 |
39 # server-side |
39 # server-side |
40 |
40 |
41 |
41 |
42 def _capabilities(orig, repo, proto): |
42 def _capabilities(orig, repo, proto): |
43 result = orig(repo, proto) |
43 result = orig(repo, proto) |
44 result.append('getannotate') |
44 result.append(b'getannotate') |
45 return result |
45 return result |
46 |
46 |
47 |
47 |
48 def _getannotate(repo, proto, path, lastnode): |
48 def _getannotate(repo, proto, path, lastnode): |
49 # output: |
49 # output: |
50 # FILE := vfspath + '\0' + str(size) + '\0' + content |
50 # FILE := vfspath + '\0' + str(size) + '\0' + content |
51 # OUTPUT := '' | FILE + OUTPUT |
51 # OUTPUT := '' | FILE + OUTPUT |
52 result = '' |
52 result = b'' |
53 buildondemand = repo.ui.configbool( |
53 buildondemand = repo.ui.configbool( |
54 'fastannotate', 'serverbuildondemand', True |
54 b'fastannotate', b'serverbuildondemand', True |
55 ) |
55 ) |
56 with context.annotatecontext(repo, path) as actx: |
56 with context.annotatecontext(repo, path) as actx: |
57 if buildondemand: |
57 if buildondemand: |
58 # update before responding to the client |
58 # update before responding to the client |
59 master = _getmaster(repo.ui) |
59 master = _getmaster(repo.ui) |
78 # agree where the main branch is. |
78 # agree where the main branch is. |
79 if actx.lastnode != lastnode: |
79 if actx.lastnode != lastnode: |
80 for p in [actx.revmappath, actx.linelogpath]: |
80 for p in [actx.revmappath, actx.linelogpath]: |
81 if not os.path.exists(p): |
81 if not os.path.exists(p): |
82 continue |
82 continue |
83 with open(p, 'rb') as f: |
83 with open(p, b'rb') as f: |
84 content = f.read() |
84 content = f.read() |
85 vfsbaselen = len(repo.vfs.base + '/') |
85 vfsbaselen = len(repo.vfs.base + b'/') |
86 relpath = p[vfsbaselen:] |
86 relpath = p[vfsbaselen:] |
87 result += '%s\0%d\0%s' % (relpath, len(content), content) |
87 result += b'%s\0%d\0%s' % (relpath, len(content), content) |
88 return result |
88 return result |
89 |
89 |
90 |
90 |
91 def _registerwireprotocommand(): |
91 def _registerwireprotocommand(): |
92 if 'getannotate' in wireprotov1server.commands: |
92 if b'getannotate' in wireprotov1server.commands: |
93 return |
93 return |
94 wireprotov1server.wireprotocommand('getannotate', 'path lastnode')( |
94 wireprotov1server.wireprotocommand(b'getannotate', b'path lastnode')( |
95 _getannotate |
95 _getannotate |
96 ) |
96 ) |
97 |
97 |
98 |
98 |
99 def serveruisetup(ui): |
99 def serveruisetup(ui): |
100 _registerwireprotocommand() |
100 _registerwireprotocommand() |
101 extensions.wrapfunction(wireprotov1server, '_capabilities', _capabilities) |
101 extensions.wrapfunction(wireprotov1server, b'_capabilities', _capabilities) |
102 |
102 |
103 |
103 |
104 # client-side |
104 # client-side |
105 |
105 |
106 |
106 |
107 def _parseresponse(payload): |
107 def _parseresponse(payload): |
108 result = {} |
108 result = {} |
109 i = 0 |
109 i = 0 |
110 l = len(payload) - 1 |
110 l = len(payload) - 1 |
111 state = 0 # 0: vfspath, 1: size |
111 state = 0 # 0: vfspath, 1: size |
112 vfspath = size = '' |
112 vfspath = size = b'' |
113 while i < l: |
113 while i < l: |
114 ch = payload[i : i + 1] |
114 ch = payload[i : i + 1] |
115 if ch == '\0': |
115 if ch == b'\0': |
116 if state == 1: |
116 if state == 1: |
117 result[vfspath] = payload[i + 1 : i + 1 + int(size)] |
117 result[vfspath] = payload[i + 1 : i + 1 + int(size)] |
118 i += int(size) |
118 i += int(size) |
119 state = 0 |
119 state = 0 |
120 vfspath = size = '' |
120 vfspath = size = b'' |
121 elif state == 0: |
121 elif state == 0: |
122 state = 1 |
122 state = 1 |
123 else: |
123 else: |
124 if state == 1: |
124 if state == 1: |
125 size += ch |
125 size += ch |
131 |
131 |
132 def peersetup(ui, peer): |
132 def peersetup(ui, peer): |
133 class fastannotatepeer(peer.__class__): |
133 class fastannotatepeer(peer.__class__): |
134 @wireprotov1peer.batchable |
134 @wireprotov1peer.batchable |
135 def getannotate(self, path, lastnode=None): |
135 def getannotate(self, path, lastnode=None): |
136 if not self.capable('getannotate'): |
136 if not self.capable(b'getannotate'): |
137 ui.warn(_('remote peer cannot provide annotate cache\n')) |
137 ui.warn(_(b'remote peer cannot provide annotate cache\n')) |
138 yield None, None |
138 yield None, None |
139 else: |
139 else: |
140 args = {'path': path, 'lastnode': lastnode or ''} |
140 args = {b'path': path, b'lastnode': lastnode or b''} |
141 f = wireprotov1peer.future() |
141 f = wireprotov1peer.future() |
142 yield args, f |
142 yield args, f |
143 yield _parseresponse(f.value) |
143 yield _parseresponse(f.value) |
144 |
144 |
145 peer.__class__ = fastannotatepeer |
145 peer.__class__ = fastannotatepeer |
148 @contextlib.contextmanager |
148 @contextlib.contextmanager |
149 def annotatepeer(repo): |
149 def annotatepeer(repo): |
150 ui = repo.ui |
150 ui = repo.ui |
151 |
151 |
152 remotepath = ui.expandpath( |
152 remotepath = ui.expandpath( |
153 ui.config('fastannotate', 'remotepath', 'default') |
153 ui.config(b'fastannotate', b'remotepath', b'default') |
154 ) |
154 ) |
155 peer = hg.peer(ui, {}, remotepath) |
155 peer = hg.peer(ui, {}, remotepath) |
156 |
156 |
157 try: |
157 try: |
158 yield peer |
158 yield peer |
173 lastnodemap = {} |
173 lastnodemap = {} |
174 |
174 |
175 ui = repo.ui |
175 ui = repo.ui |
176 results = [] |
176 results = [] |
177 with peer.commandexecutor() as batcher: |
177 with peer.commandexecutor() as batcher: |
178 ui.debug('fastannotate: requesting %d files\n' % len(paths)) |
178 ui.debug(b'fastannotate: requesting %d files\n' % len(paths)) |
179 for p in paths: |
179 for p in paths: |
180 results.append( |
180 results.append( |
181 batcher.callcommand( |
181 batcher.callcommand( |
182 'getannotate', {'path': p, 'lastnode': lastnodemap.get(p)} |
182 b'getannotate', |
|
183 {b'path': p, b'lastnode': lastnodemap.get(p)}, |
183 ) |
184 ) |
184 ) |
185 ) |
185 |
186 |
186 for result in results: |
187 for result in results: |
187 r = result.result() |
188 r = result.result() |
188 # TODO: pconvert these paths on the server? |
189 # TODO: pconvert these paths on the server? |
189 r = {util.pconvert(p): v for p, v in r.iteritems()} |
190 r = {util.pconvert(p): v for p, v in r.iteritems()} |
190 for path in sorted(r): |
191 for path in sorted(r): |
191 # ignore malicious paths |
192 # ignore malicious paths |
192 if not path.startswith('fastannotate/') or '/../' in ( |
193 if not path.startswith(b'fastannotate/') or b'/../' in ( |
193 path + '/' |
194 path + b'/' |
194 ): |
195 ): |
195 ui.debug('fastannotate: ignored malicious path %s\n' % path) |
196 ui.debug( |
|
197 b'fastannotate: ignored malicious path %s\n' % path |
|
198 ) |
196 continue |
199 continue |
197 content = r[path] |
200 content = r[path] |
198 if ui.debugflag: |
201 if ui.debugflag: |
199 ui.debug( |
202 ui.debug( |
200 'fastannotate: writing %d bytes to %s\n' |
203 b'fastannotate: writing %d bytes to %s\n' |
201 % (len(content), path) |
204 % (len(content), path) |
202 ) |
205 ) |
203 repo.vfs.makedirs(os.path.dirname(path)) |
206 repo.vfs.makedirs(os.path.dirname(path)) |
204 with repo.vfs(path, 'wb') as f: |
207 with repo.vfs(path, b'wb') as f: |
205 f.write(content) |
208 f.write(content) |
206 |
209 |
207 |
210 |
208 def _filterfetchpaths(repo, paths): |
211 def _filterfetchpaths(repo, paths): |
209 """return a subset of paths whose history is long and need to fetch linelog |
212 """return a subset of paths whose history is long and need to fetch linelog |
210 from the server. works with remotefilelog and non-remotefilelog repos. |
213 from the server. works with remotefilelog and non-remotefilelog repos. |
211 """ |
214 """ |
212 threshold = repo.ui.configint('fastannotate', 'clientfetchthreshold', 10) |
215 threshold = repo.ui.configint(b'fastannotate', b'clientfetchthreshold', 10) |
213 if threshold <= 0: |
216 if threshold <= 0: |
214 return paths |
217 return paths |
215 |
218 |
216 result = [] |
219 result = [] |
217 for path in paths: |
220 for path in paths: |
238 lastnodemap[path] = actx.lastnode |
241 lastnodemap[path] = actx.lastnode |
239 if needupdatepaths: |
242 if needupdatepaths: |
240 clientfetch(self, needupdatepaths, lastnodemap, peer) |
243 clientfetch(self, needupdatepaths, lastnodemap, peer) |
241 except Exception as ex: |
244 except Exception as ex: |
242 # could be directory not writable or so, not fatal |
245 # could be directory not writable or so, not fatal |
243 self.ui.debug('fastannotate: prefetch failed: %r\n' % ex) |
246 self.ui.debug(b'fastannotate: prefetch failed: %r\n' % ex) |
244 |
247 |
245 repo.__class__ = fastannotaterepo |
248 repo.__class__ = fastannotaterepo |
246 |
249 |
247 |
250 |
248 def clientreposetup(ui, repo): |
251 def clientreposetup(ui, repo): |