1#!/usr/bin/env python3
2"""Compare SMB lock progress and exclusion across file access modes."""
3import argparse
4from collections import Counter
5import fcntl
6import json
7import hashlib
8import platform
9import os
10from pathlib import Path
11import subprocess
12import sys
13import time
14
15
16def worker(file, marker, mode, operations, atomic, shared):
17 counts = Counter()
18 for i in range(operations):
19 access = ('r+b' if i % 2 else 'rb') if mode == 'mixed' else mode
20 # With shared reads, readers open O_SHLOCK and may overlap each other, never a writer.
21 reading = shared and access == 'rb'
22 held = marker.with_name(f'{marker.name}-reader-{os.getpid()}') if reading else marker
23 stage = 'open'
24 try:
25 lock = os.O_SHLOCK if reading else os.O_EXLOCK
26 with (os.fdopen(os.open(file, (os.O_RDWR if access == 'r+b' else os.O_RDONLY) | lock | os.O_NONBLOCK), access) if atomic else file.open(access)) as stream:
27 stage = 'lock'
28 if not atomic:
29 fcntl.flock(stream, fcntl.LOCK_EX | fcntl.LOCK_NB)
30 counts['acquired'] += 1
31 owned = False
32 try:
33 try:
34 held.mkdir()
35 owned = True
36 others = [p for p in marker.parent.glob(marker.name + '*') if p != held]
37 if others and not (reading and all('-reader-' in p.name for p in others)):
38 counts['overlap'] += 1
39 elif others:
40 counts['shared'] += 1
41 except FileExistsError:
42 counts['overlap'] += 1
43 time.sleep(.002)
44 try:
45 stream.read(1)
46 except OSError as error:
47 counts[f'read:{error.errno}'] += 1
48 finally:
49 if owned:
50 held.rmdir()
51 try:
52 fcntl.flock(stream, fcntl.LOCK_UN)
53 except OSError as error:
54 counts[f'unlock:{error.errno}'] += 1
55 stage = 'close'
56 except OSError as error:
57 counts[f'{stage}:{error.errno}'] += 1
58 time.sleep(.002)
59 print(json.dumps(counts), flush=True)
60
61
62def run(output, directory, operations, atomic, clients, shared):
63 output.mkdir(parents=True, exist_ok=False)
64 directory.mkdir(exist_ok=False)
65 (output / 'run.json').write_text(json.dumps({'operations': operations, 'clients': clients, 'atomic': atomic, 'shared_reads': shared,
66 'platform': platform.platform(), 'script_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}, indent=2))
67 results = {}
68 for mode in ('rb', 'r+b', 'mixed'):
69 file = directory / (mode + '.bin')
70 file.write_bytes(b'lock probe')
71 marker = output / 'held'
72 processes = []
73 try:
74 for i in range(clients):
75 processes.append(subprocess.Popen([sys.executable, __file__, '--worker', str(file), str(marker), mode, str(operations), str(int(atomic)), str(int(shared))],
76 stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True))
77 results[mode] = []
78 for process in processes:
79 stdout, stderr = process.communicate(timeout=60)
80 assert process.returncode == 0, stderr
81 results[mode].append(json.loads(stdout))
82 finally:
83 for process in processes:
84 if process.poll() is None:
85 process.terminate()
86 for process in processes:
87 try:
88 process.wait(timeout=5)
89 except subprocess.TimeoutExpired:
90 process.kill()
91 process.wait()
92 (output / 'result.json').write_text(json.dumps(results, indent=2))
93 print(mode, results[mode], flush=True)
94 assert all(c.get('acquired', 0) and not c.get('overlap', 0) and
95 all(k in ('acquired', 'shared', 'open:35', 'lock:35') for k in c)
96 for counters in results.values() for c in counters), 'Lock exclusion or progress failed'
97
98
99if __name__ == '__main__':
100 if len(sys.argv) == 8 and sys.argv[1] == '--worker':
101 worker(Path(sys.argv[2]), Path(sys.argv[3]), sys.argv[4], int(sys.argv[5]), bool(int(sys.argv[6])), bool(int(sys.argv[7])))
102 else:
103 parser = argparse.ArgumentParser(description=__doc__)
104 parser.add_argument('output', type=Path)
105 parser.add_argument('directory', type=Path)
106 parser.add_argument('--operations', type=int, default=500)
107 parser.add_argument('--atomic', action='store_true')
108 parser.add_argument('--clients', type=int, default=4)
109 parser.add_argument('--shared-reads', action='store_true', help='With --atomic, readers open O_SHLOCK.')
110 args = parser.parse_args()
111 run(args.output.resolve(), args.directory.resolve(), args.operations, args.atomic, args.clients, args.shared_reads)