io_v2_tcp.py (14901B)
1 #!/usr/bin/env python3 2 """Real socket source workflows: full binary echo, directional pipelines and handoff.""" 3 import argparse 4 import concurrent.futures 5 import hashlib 6 import json 7 from pathlib import Path 8 import socket 9 import subprocess 10 import struct 11 import threading 12 import time 13 from source_order import Runner 14 from process_rss import ChildRSS 15 from rts_metrics import parse_rts_stats 16 from http_headers import number, result, tree_list 17 from io_v2_positioned import text 18 from io_v2_directory import parse_tree, items, string 19 ROOT = Path(__file__).resolve().parent.parent 20 21 def endpoint(family, port): 22 if family == 4: 23 return tree_list([text('ipv4'),tree_list([number(x) for x in [127,0,0,1]]),number(port)]) 24 return tree_list([text('ipv6'),tree_list([number(x) for x in [0,0,0,0,0,0,0,1]]),number(0),number(port)]) 25 26 def exchange(family, port, body, prefix): 27 with socket.socket(socket.AF_INET if family == 4 else socket.AF_INET6) as peer: 28 peer.settimeout(90) 29 peer.connect(('127.0.0.1' if family == 4 else '::1',port)) 30 def send(): 31 for start in range(0,len(body),997): 32 peer.sendall(body[start:start+997]) 33 peer.shutdown(socket.SHUT_WR) 34 with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: 35 writer = pool.submit(send) 36 chunks = [] 37 while True: 38 chunk = peer.recv(4096) 39 if not chunk: 40 break 41 chunks.append(chunk) 42 writer.result(timeout=90) 43 actual = b''.join(chunks) 44 assert actual == prefix+body, (len(actual),len(prefix+body)) 45 return hashlib.sha256(actual).hexdigest() 46 47 def main(): 48 parser = argparse.ArgumentParser(description=__doc__) 49 parser.add_argument('tricu', type=Path) 50 parser.add_argument('--directory', type=Path, required=True) 51 args = parser.parse_args() 52 directory = args.directory.resolve() 53 directory.mkdir(parents=True, exist_ok=False) 54 runner = Runner(args.tricu,directory) 55 source = ROOT/'demos/io-v2/tcp.tri' 56 bundle = directory/'tcp.arboricx' 57 runner.run('compile',['compile','--file',source,'--store',runner.store,'--output',bundle]) 58 pin = runner.imported(bundle,'import') 59 checks, observations = [], [] 60 def command(label, name, values): 61 argv = [str(runner.exe),'arboricx','run','--module',pin,'--name',name,'--store',str(runner.store),'--io','--no-result'] 62 for index,value in enumerate(values): 63 path = directory/f'{label}-{index}.tree' 64 path.write_bytes(value) 65 argv += ['--arg-tree-file',str(path)] 66 return argv+['+RTS','-N1','-M128m','-T','-s','-RTS'] 67 def launch(label,name,values): 68 argv = command(label,name,values) 69 proc = subprocess.Popen(argv,stdin=subprocess.DEVNULL,stdout=subprocess.PIPE,stderr=subprocess.PIPE) 70 data = dict(label=label, command=argv, started=time.monotonic(), peak_kib=0) 71 stopped = threading.Event() 72 def sample(): 73 try: 74 rss = ChildRSS(proc) 75 while not stopped.is_set(): 76 current = rss.sample() 77 if current is not None: 78 data['peak_kib'] = max(data['peak_kib'],current) 79 if current > 128*1024: 80 raise AssertionError('RSS bound exceeded') 81 stopped.wait(.01) 82 except Exception as error: 83 data['monitor_error'] = repr(error) 84 try: 85 proc.kill() 86 except ProcessLookupError: 87 pass 88 thread = threading.Thread(target=sample,daemon=True) 89 thread.start() 90 return proc,data,stopped,thread 91 def finish(run): 92 proc,data,stopped,thread = run 93 try: 94 stdout,stderr = proc.communicate(timeout=120) 95 stopped.set(); thread.join() 96 assert 'monitor_error' not in data, data 97 assert proc.returncode == 0, stderr 98 assert not stdout, stdout 99 stats = parse_rts_stats(stderr.decode()) 100 assert stats['maximum_live_bytes'] <= 16*1024*1024, stats 101 assert 0 < data['peak_kib'] <= 128*1024, data 102 (directory/(data['label']+'.log')).write_bytes(stderr) 103 data['seconds'] = time.monotonic()-data.pop('started') 104 data['rts_process'] = stats 105 observations.append(data) 106 finally: 107 stopped.set(); thread.join() 108 if proc.poll() is None: 109 proc.kill(); proc.communicate() 110 def abort(run): 111 proc,_,stopped,thread = run 112 stopped.set(); thread.join() 113 if proc.poll() is None: 114 proc.kill(); proc.communicate() 115 def server(label, name, family, bodies, parallel=False): 116 run = launch(label,name,[endpoint(family,0),number(len(bodies))]) 117 proc = run[0] 118 try: 119 with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: 120 # A timed readiness channel, never a probe connection consuming 121 # one of the finite server's accepted-connection slots. 122 ready = pool.submit(proc.stdout.readline) 123 try: 124 line = ready.result(timeout=20) 125 if not line: 126 _, diagnostic = proc.communicate(timeout=5) 127 (directory/(label+'.startup-error.log')).write_bytes(diagnostic) 128 raise AssertionError(diagnostic) 129 port = int(line) 130 except BaseException: 131 abort(run) 132 raise 133 prefix = bytes([66,0,255,10]) if name == 'handledServer' else b'' 134 if parallel: 135 with concurrent.futures.ThreadPoolExecutor(max_workers=len(bodies)) as pool: 136 sums = list(pool.map(lambda body: exchange(family,port,body,prefix),bodies)) 137 else: 138 sums = [exchange(family,port,body,prefix) for body in bodies] 139 finish(run) 140 checks.append(dict(label=label,sha256=sums)) 141 finally: 142 abort(run) 143 body = lambda size: (bytes(range(256))*((size+255)//256))[:size] 144 for name in ['server','handledServer']: 145 server(name+'-boundaries',name,4,[body(n) for n in [0,1,4095,4096,4097,65536,1048576]]) 146 server(name+'-ipv6',name,6,[body(4097)]) 147 server(name+'-concurrent',name,4,[bytes([i])+body(65535) for i in range(4)],True) 148 server(name+'-churn',name,4,[bytes([i])*64 for i in range(128)]) 149 for name in ['handledServer','supervisedHandledServer']: 150 label = name+'-reset' 151 run = launch(label,name,[endpoint(4,0),number(2)]) 152 proc = run[0] 153 try: 154 with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: 155 ready = pool.submit(proc.stdout.readline) 156 try: 157 port = int(ready.result(timeout=20)) 158 except BaseException: 159 abort(run) 160 raise 161 with socket.create_connection(('127.0.0.1',port),timeout=20) as peer: 162 banner = b'' 163 while len(banner) < 4: 164 chunk = peer.recv(4-len(banner)) 165 assert chunk, banner 166 banner += chunk 167 assert banner == bytes([66,0,255,10]) 168 # Banner proves accept/handoff completed before aborting the peer. 169 peer.setsockopt(socket.SOL_SOCKET,socket.SO_LINGER,struct.pack('ii',1,0)) 170 if name == 'supervisedHandledServer': 171 exchange(4,port,body(4097),bytes([66,0,255,10])) 172 finish(run) 173 else: 174 stdout,stderr = proc.communicate(timeout=20) 175 assert proc.returncode != 0 and b'IO v2 task failed' in stderr, stderr 176 assert stdout == b'' 177 (directory/(label+'.log')).write_bytes(stderr) 178 checks.append(dict(label=label)) 179 finally: 180 abort(run) 181 # Native source client reads/writes its socket concurrently. A real peer 182 # echoes with small buffers, forcing meaningful bidirectional backpressure. 183 for family in [4,6]: 184 for size in ([0,4097,1048576] if family == 4 else [4097]): 185 label = f'client-{family}-{size}' 186 incoming, outgoing = directory/(label+'.in'), directory/(label+'.out') 187 incoming.write_bytes(body(size)) 188 with socket.socket(socket.AF_INET if family == 4 else socket.AF_INET6) as listener: 189 listener.bind(('127.0.0.1' if family == 4 else '::1',0)); listener.listen(1); listener.settimeout(20) 190 port = listener.getsockname()[1] 191 def echo(): 192 peer,_ = listener.accept() 193 with peer: 194 peer.settimeout(90) 195 while True: 196 chunk = peer.recv(997) 197 if not chunk: 198 break 199 peer.sendall(chunk) 200 peer.shutdown(socket.SHUT_WR) 201 with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: 202 worker = pool.submit(echo) 203 run = launch(label,'client',[endpoint(family,port),text(str(incoming)),text(str(outgoing))]) 204 try: 205 finish(run) 206 worker.result(timeout=90) 207 finally: 208 abort(run) 209 assert outgoing.read_bytes() == incoming.read_bytes(), label 210 checks.append(dict(label=label,sha256=hashlib.sha256(outgoing.read_bytes()).hexdigest())) 211 # Explicit host resolution plus source fallback, using only local services. 212 for host in ['127.0.0.1','localhost']: 213 label = 'named-client-'+host 214 incoming, outgoing = directory/(label+'.in'), directory/(label+'.out') 215 incoming.write_bytes(body(65536)) 216 with socket.socket() as listener: 217 listener.bind(('127.0.0.1',0)); listener.listen(1); listener.settimeout(20) 218 port = listener.getsockname()[1] 219 def echo_named(): 220 peer,_ = listener.accept() 221 with peer: 222 peer.settimeout(30) 223 while True: 224 chunk = peer.recv(997) 225 if not chunk: 226 break 227 peer.sendall(chunk) 228 peer.shutdown(socket.SHUT_WR) 229 with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: 230 peer = pool.submit(echo_named) 231 run = launch(label,'namedClient',[text(host),text('any'),number(port),text(str(incoming)),text(str(outgoing))]) 232 try: 233 finish(run) 234 peer.result(timeout=30) 235 finally: 236 abort(run) 237 assert outgoing.read_bytes() == incoming.read_bytes() 238 checks.append(dict(label=label,sha256=hashlib.sha256(outgoing.read_bytes()).hexdigest())) 239 def invoke_values(label,name,values): 240 argv = ['--io'] 241 for index,value in enumerate(values): 242 path = directory/f'{label}-{index}.tree'; path.write_bytes(value) 243 argv += ['--arg-tree-file',path] 244 return runner.invoke(label,pin,name,argv) 245 actual = parse_tree(invoke_values('resolve-numeric','resolve',[text('127.0.0.1'),text('ipv4'),number(1234),number(64)])) 246 wrapped = items(actual) 247 assert wrapped[0] == (1,(0,)) 248 addresses,truncated = items(wrapped[1]) 249 assert items(addresses) and all(x == parse_tree(endpoint(4,1234)) for x in items(addresses)) 250 assert truncated == (0,) 251 checks.append(dict(label='resolve-numeric')) 252 for label,host,limit,category,detail in [ 253 ('resolve-host','bad/name',1,'invalid-argument','expected bounded ASCII host name'), 254 ('resolve-limit','localhost',0,'invalid-bound',''), 255 ]: 256 expected_error = result(False,tree_list([text(x) for x in [category,'resolver/resolve',detail,'none']])) 257 assert invoke_values(label,'resolve',[text(host),text('any'),number(80),number(limit)]) == expected_error 258 checks.append(dict(label=label)) 259 with socket.socket() as listener: 260 listener.bind(('127.0.0.1',0)); listener.listen(1); listener.settimeout(20) 261 selected = listener.getsockname()[1] 262 # Invalid first candidate, valid second, malformed third. Source policy 263 # retries ordinary failure but must stop before inspecting the third. 264 reply = items(parse_tree(invoke_values('candidate-fallback','candidates', 265 [tree_list([endpoint(4,0),endpoint(4,selected),b'\1\0'])]))) 266 assert reply[0] == (1,(0,)) 267 token = items(reply[1]) 268 assert len(token) == 3 and string(token[0]) == 'io.v2.resource' 269 peer,_ = listener.accept() 270 with peer: 271 peer.settimeout(20) 272 assert peer.recv(1) == b'' # returned token did not extend its scope 273 checks.append(dict(label='candidate-fallback')) 274 empty_error = result(False,tree_list([text(x) for x in ['not-found','tcp/connect-first','No endpoint candidates','none']])) 275 assert invoke_values('empty-candidates','candidates',[b'\0']) == empty_error 276 checks.append(dict(label='empty-candidates')) 277 expected = result(False,tree_list([text(x) for x in ['invalid-argument','tcp/connect','','none']])) 278 assert runner.invoke('invalid-address',pin,'invalidAddress',['--io']) == expected 279 checks.append(dict(label='invalid-address')) 280 with socket.socket() as unused: 281 unused.bind(('127.0.0.1',0)) 282 port = unused.getsockname()[1] 283 path = directory/'refused.tree'; path.write_bytes(endpoint(4,port)) 284 reply = items(parse_tree(runner.invoke('refused',pin,'connectResult',['--io','--arg-tree-file',path]))) 285 assert reply[0] == (0,) 286 problem = items(reply[1]) 287 assert [string(problem[i]) for i in [0,1,3]] == ['connection-refused','tcp/connect','unknown'] 288 checks.append(dict(label='refused')) 289 paths = [source,Path(__file__).resolve(),ROOT/'lib/io-v2.tri',ROOT/'lib/handlers.tri', 290 ROOT/'src/IOV2/Network.hs',ROOT/'src/IOV2/Resolver.hs',ROOT/'src/IOV2/Resources.hs',ROOT/'src/IOV2/Driver.hs', 291 ROOT/'src/IOV2/Protocol.hs',ROOT/'test/rts_metrics.py',ROOT/'test/process_rss.py',ROOT/'bootstrap/build-stable.arboricx'] 292 (directory/'report.json').write_text(json.dumps(dict(checks=checks,observations=observations, 293 sources={str(p.relative_to(ROOT)):hashlib.sha256(p.read_bytes()).hexdigest() for p in paths}),indent=2)+'\n') 294 print(f'PASS: {len(checks)} actual TCP workflow groups') 295 if __name__ == '__main__': 296 main()