package rpc_parallel

  1. Overview
  2. Docs
Type-safe parallel library built on top of Async_rpc

Install

Dune Dependency

Authors

Maintainers

Sources

v0.17.0.tar.gz
sha256=160c3c60b224f3238810858435e8ce5d51376edf6fe2af6cc0ed02edf0166e08

doc/src/rpc_parallel/parallel.ml.html

Source file parallel.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
open Core
open Async
open Parallel_intf
module Worker_type_id = Utils.Worker_type_id
module Worker_id = Utils.Worker_id

(* All processes start a "master" rpc server. This is a server that has two
   implementations:

   (1) Register - Spawned workers say hello to the spawner
   (2) Handle_exn - Spawned workers send exceptions to the spawner

   Processes can also start "worker" rpc servers (in process using [serve] or out of
   process using [spawn]). A "worker" rpc server has all the user defined implementations
   as well as:

   (1) Init_worker_state - Spawner sends the worker init argument
   (2) Init_connection_state - Connector sends the connection state init argument
   (3) Shutdown, Close_server, Async_log, etc.

   The handshake protocol for spawning a worker:

   (Master) - SSH and start running executable
   (Worker) - Start server, send [Register_rpc] with its host and port
   (Master) - Connect to worker, send [Init_worker_state_rpc]
   (Worker) - Do initialization (ensure we have daemonized first)
   (Master) - Finally, return a [Worker.t] to the caller
*)

module Worker_implementations = struct
  type t =
    | T :
        ('state, 'connection_state) Utils.Internal_connection_state.t Rpc.Implementation.t
        list
        -> t
end

(* Applications of the [Make()] functor have the side effect of populating an
   [implementations] list which subsequently adds an entry for that worker type id to the
   [worker_implementations]. *)
let worker_implementations = Worker_type_id.Table.create ~size:1 ()

module Worker_command_args = struct
  type t =
    | Add_master_pid
    | User_supplied of
        { args : string list
        ; pass_name : bool
        }
  [@@deriving sexp]
end

module Server_with_rpc_settings = struct
  type t =
    { server : (Socket.Address.Inet.t, int) Tcp.Server.t
    ; rpc_settings : Rpc_settings.t
    }
end

(* Functions that are implemented by all workers *)
module Shutdown_rpc = struct
  let rpc = Rpc.One_way.create ~name:"shutdown_rpc" ~version:0 ~bin_msg:Unit.bin_t
end

module Close_server_rpc = struct
  let rpc = Rpc.One_way.create ~name:"close_server_rpc" ~version:0 ~bin_msg:Unit.bin_t
end

module Worker_server_rpc_settings_rpc = struct
  let rpc =
    Rpc.Rpc.create
      ~name:"worker_server_rpc_settings_rpc"
      ~version:0
      ~bin_query:Unit.bin_t
      ~bin_response:Rpc_settings.bin_t
      ~include_in_error_count:Only_on_exn
  ;;
end

module Async_log_rpc = struct
  let rpc =
    Rpc.Pipe_rpc.create
      ~name:"async_log_rpc"
      ~version:0
      ~bin_query:Unit.bin_t
      ~bin_response:Log.Message.Stable.V2.bin_t
      ~bin_error:Error.bin_t
      ()
  ;;
end

module Function = struct
  module Rpc_id = Unique_id.Int ()

  let maybe_generate_name ~prefix ~name =
    match name with
    | None -> sprintf "%s_%s" prefix (Rpc_id.to_string (Rpc_id.create ()))
    | Some n -> n
  ;;

  module Function_piped = struct
    type ('worker, 'query, 'response) t = ('query, 'response, Error.t) Rpc.Pipe_rpc.t

    let make_impl ~monitor ~f protocol =
      Rpc.Pipe_rpc.implement
        protocol
        (fun ((_conn : Rpc.Connection.t), internal_conn_state) arg ->
        let { Utils.Internal_connection_state.conn_state; worker_state; _ } =
          Set_once.get_exn internal_conn_state [%here]
        in
        Utils.try_within ~monitor (fun () -> f ~worker_state ~conn_state arg))
    ;;

    let make_direct_impl ~monitor ~f protocol =
      Rpc.Pipe_rpc.implement_direct
        protocol
        (fun ((_conn : Rpc.Connection.t), internal_conn_state) arg writer ->
        let { Utils.Internal_connection_state.conn_state; worker_state; _ } =
          Set_once.get_exn internal_conn_state [%here]
        in
        Utils.try_within ~monitor (fun () -> f ~worker_state ~conn_state arg writer))
    ;;

    let make_proto ~name ~bin_input ~bin_output ~client_pushes_back =
      let name = maybe_generate_name ~prefix:"rpc_parallel_piped" ~name in
      Rpc.Pipe_rpc.create
        ?client_pushes_back
        ~name
        ~version:0
        ~bin_query:bin_input
        ~bin_response:bin_output
        ~bin_error:Error.bin_t
        ()
    ;;
  end

  module Function_state = struct
    type ('worker, 'query, 'state, 'update) t =
      ('query, 'state, 'update, Error.t) Rpc.State_rpc.t

    let make_impl ~monitor ~f protocol =
      Rpc.State_rpc.implement
        protocol
        (fun ((_conn : Rpc.Connection.t), internal_conn_state) arg ->
        let { Utils.Internal_connection_state.conn_state; worker_state; _ } =
          Set_once.get_exn internal_conn_state [%here]
        in
        Utils.try_within ~monitor (fun () -> f ~worker_state ~conn_state arg))
    ;;

    let make_proto ~name ~bin_query ~bin_state ~bin_update ~client_pushes_back =
      let name = maybe_generate_name ~prefix:"rpc_parallel_state" ~name in
      Rpc.State_rpc.create
        ?client_pushes_back
        ~name
        ~version:0
        ~bin_query
        ~bin_state
        ~bin_update
        ~bin_error:Error.bin_t
        ()
    ;;
  end

  module Function_plain = struct
    type ('worker, 'query, 'response) t = ('query, 'response) Rpc.Rpc.t

    let make_impl ~monitor ~f protocol =
      Rpc.Rpc.implement
        protocol
        (fun ((_conn : Rpc.Connection.t), internal_conn_state) arg ->
        let { Utils.Internal_connection_state.conn_state; worker_state; _ } =
          Set_once.get_exn internal_conn_state [%here]
        in
        (* We want to raise any exceptions from [f arg] to the current monitor (handled
              by Rpc) so the caller can see it. Additional exceptions will be handled by the
              specified monitor *)
        Utils.try_within_exn ~monitor (fun () -> f ~worker_state ~conn_state arg))
    ;;

    let make_proto ~name ~bin_input ~bin_output =
      let name = maybe_generate_name ~prefix:"rpc_parallel_plain" ~name in
      Rpc.Rpc.create
        ~name
        ~version:0
        ~bin_query:bin_input
        ~bin_response:bin_output
        ~include_in_error_count:Only_on_exn
    ;;
  end

  module Function_reverse_piped = struct
    module Id = Unique_id.Int ()

    type ('worker, 'query, 'update, 'response, 'in_progress) t =
      { worker_rpc : ('query * Id.t, 'response) Rpc.Rpc.t
      ; master_rpc : (Id.t, 'update, Error.t) Rpc.Pipe_rpc.t
      ; master_in_progress : 'in_progress Id.Table.t
      }

    let make_worker_impl ~monitor ~f t =
      Rpc.Rpc.implement t.worker_rpc (fun (conn, internal_conn_state) (arg, id) ->
        let { Utils.Internal_connection_state.conn_state; worker_state; _ } =
          Set_once.get_exn internal_conn_state [%here]
        in
        Utils.try_within_exn ~monitor (fun () ->
          match%bind Rpc.Pipe_rpc.dispatch t.master_rpc conn id with
          | Ok (Ok (updates, (_ : Rpc.Pipe_rpc.Metadata.t))) ->
            f ~worker_state ~conn_state arg updates
          | Ok (Error error) | Error error -> Error.raise error))
    ;;

    let make_master t ~implement ~ok ~error =
      implement t.master_rpc (fun () id ->
        match Hashtbl.find_and_remove t.master_in_progress id with
        | Some in_progress -> ok in_progress
        | None ->
          [%message
            "Bug in Rpc_parallel: reverse pipe master implementation not found"
              (id : Id.t)
              (Rpc.Pipe_rpc.name t.master_rpc : string)]
          |> Deferred.Or_error.error_s
          |> error)
    ;;

    let make_master_impl t =
      make_master
        t
        ~implement:Rpc.Pipe_rpc.implement
        ~ok:Deferred.Or_error.return
        ~error:Fn.id
    ;;

    let make_master_impl_direct t =
      make_master t ~implement:Rpc.Pipe_rpc.implement_direct ~ok:Fn.id ~error:const
    ;;

    let make_proto ~name ~bin_query ~bin_update ~bin_response ~client_pushes_back =
      let name = maybe_generate_name ~prefix:"rpc_parallel_reverse_piped" ~name in
      let worker_rpc =
        let module With_id = struct
          type 'a t = 'a * Id.t [@@deriving bin_io]
        end
        in
        Rpc.Rpc.create
          ~name
          ~version:0
          ~bin_query:(With_id.bin_t bin_query)
          ~bin_response
          ~include_in_error_count:Only_on_exn
      in
      let master_rpc =
        Rpc.Pipe_rpc.create
          ?client_pushes_back
          ~name
          ~version:0
          ~bin_query:Id.bin_t
          ~bin_response:bin_update
          ~bin_error:Error.bin_t
          ()
      in
      let master_in_progress = Id.Table.create () in
      { worker_rpc; master_rpc; master_in_progress }
    ;;
  end

  module Function_one_way = struct
    type ('worker, 'query) t = 'query Rpc.One_way.t

    let make_impl ~monitor ~f protocol =
      Rpc.One_way.implement
        protocol
        (fun ((_conn : Rpc.Connection.t), internal_conn_state) arg ->
        let { Utils.Internal_connection_state.conn_state; worker_state; _ } =
          Set_once.get_exn internal_conn_state [%here]
        in
        don't_wait_for
          (* Even though [f] returns [unit], we want to use [try_within_exn] so if it
                starts any background jobs we won't miss the exceptions *)
          (Utils.try_within_exn ~monitor (fun () ->
             f ~worker_state ~conn_state arg |> return)))
    ;;

    let make_proto ~name ~bin_input =
      let name =
        match name with
        | None -> sprintf "rpc_parallel_one_way_%s" (Rpc_id.to_string (Rpc_id.create ()))
        | Some n -> n
      in
      Rpc.One_way.create ~name ~version:0 ~bin_msg:bin_input
    ;;
  end

  module Id_directly_piped = struct
    type 'worker t = T : _ Rpc.Pipe_rpc.t * Rpc.Pipe_rpc.Id.t -> 'worker t

    let abort (T (proto, id)) connection = Rpc.Pipe_rpc.abort proto connection id
  end

  type ('worker, 'query, 'response) t_internal =
    | Plain of ('worker, 'query, 'response) Function_plain.t
    | Piped :
        ('worker, 'query, 'response) Function_piped.t
        * ('r, 'response Pipe.Reader.t) Type_equal.t
        -> ('worker, 'query, 'r) t_internal
    | State :
        ('worker, 'query, 'state, 'update) Function_state.t
        * ('r, 'state * 'update Pipe.Reader.t) Type_equal.t
        -> ('worker, 'query, 'r) t_internal
    | Directly_piped :
        ('worker, 'query, 'response) Function_piped.t
        -> ( 'worker
           , 'query
             * ('response Rpc.Pipe_rpc.Pipe_message.t -> Rpc.Pipe_rpc.Pipe_response.t)
           , 'worker Id_directly_piped.t )
           t_internal
    | One_way : ('worker, 'query) Function_one_way.t -> ('worker, 'query, unit) t_internal
    | Reverse_piped :
        ('worker, 'query, 'update, 'response, 'in_progress) Function_reverse_piped.t
        * ('q, 'query * 'in_progress) Type_equal.t
        -> ('worker, 'q, 'response) t_internal

  type ('worker, 'query, +'response) t =
    | T :
        ('query -> 'query_internal)
        * ('worker, 'query_internal, 'response_internal) t_internal
        * ('response_internal -> 'response)
        -> ('worker, 'query, 'response) t

  module Direct_pipe = struct
    module Id = Id_directly_piped

    type nonrec ('worker, 'query, 'response) t =
      ( 'worker
      , 'query * ('response Rpc.Pipe_rpc.Pipe_message.t -> Rpc.Pipe_rpc.Pipe_response.t)
      , 'worker Id.t )
      t
  end

  let map (T (q, i, r)) ~f = T (q, i, Fn.compose f r)
  let contra_map (T (q, i, r)) ~f = T (Fn.compose q f, i, r)

  let create_rpc ~monitor ~name ~f ~bin_input ~bin_output =
    let proto = Function_plain.make_proto ~name ~bin_input ~bin_output in
    let impl = Function_plain.make_impl ~monitor ~f proto in
    T (Fn.id, Plain proto, Fn.id), impl
  ;;

  let create_pipe ~monitor ~name ~f ~bin_input ~bin_output ~client_pushes_back =
    let proto =
      Function_piped.make_proto ~name ~bin_input ~bin_output ~client_pushes_back
    in
    let impl = Function_piped.make_impl ~monitor ~f proto in
    T (Fn.id, Piped (proto, Type_equal.T), Fn.id), impl
  ;;

  let create_state ~monitor ~name ~f ~bin_query ~bin_state ~bin_update ~client_pushes_back
    =
    let proto =
      Function_state.make_proto
        ~name
        ~bin_query
        ~bin_state
        ~bin_update
        ~client_pushes_back
    in
    let impl = Function_state.make_impl ~monitor ~f proto in
    T (Fn.id, State (proto, Type_equal.T), Fn.id), impl
  ;;

  let create_direct_pipe ~monitor ~name ~f ~bin_input ~bin_output ~client_pushes_back =
    let proto =
      Function_piped.make_proto ~name ~bin_input ~bin_output ~client_pushes_back
    in
    let impl = Function_piped.make_direct_impl ~monitor ~f proto in
    T (Fn.id, Directly_piped proto, Fn.id), impl
  ;;

  let create_one_way ~monitor ~name ~f ~bin_input =
    let proto = Function_one_way.make_proto ~name ~bin_input in
    let impl = Function_one_way.make_impl ~monitor ~f proto in
    T (Fn.id, One_way proto, Fn.id), impl
  ;;

  let reverse_pipe
    ~make_master_impl
    ~monitor
    ~name
    ~f
    ~bin_query
    ~bin_update
    ~bin_response
    ~client_pushes_back
    =
    let proto =
      Function_reverse_piped.make_proto
        ~name
        ~bin_query
        ~bin_update
        ~bin_response
        ~client_pushes_back
    in
    let worker_impl = Function_reverse_piped.make_worker_impl ~monitor ~f proto in
    let master_impl = make_master_impl proto in
    ( T (Fn.id, Reverse_piped (proto, Type_equal.T), Fn.id)
    , `Worker worker_impl
    , `Master master_impl )
  ;;

  let create_reverse_pipe ~monitor ~name ~f ~bin_query ~bin_update ~bin_response =
    reverse_pipe
      ~make_master_impl:Function_reverse_piped.make_master_impl
      ~monitor
      ~name
      ~f
      ~bin_query
      ~bin_update
      ~bin_response
  ;;

  let create_reverse_direct_pipe ~monitor ~name ~f ~bin_query ~bin_update ~bin_response =
    reverse_pipe
      ~make_master_impl:Function_reverse_piped.make_master_impl_direct
      ~monitor
      ~name
      ~f
      ~bin_query
      ~bin_update
      ~bin_response
  ;;

  let of_async_rpc ~monitor ~f proto =
    let impl = Function_plain.make_impl ~monitor ~f proto in
    T (Fn.id, Plain proto, Fn.id), impl
  ;;

  let of_async_pipe_rpc ~monitor ~f proto =
    let impl = Function_piped.make_impl ~monitor ~f proto in
    T (Fn.id, Piped (proto, Type_equal.T), Fn.id), impl
  ;;

  let of_async_state_rpc ~monitor ~f proto =
    let impl = Function_state.make_impl ~monitor ~f proto in
    T (Fn.id, State (proto, Type_equal.T), Fn.id), impl
  ;;

  let of_async_direct_pipe_rpc ~monitor ~f proto =
    let impl = Function_piped.make_direct_impl ~monitor ~f proto in
    T (Fn.id, Directly_piped proto, Fn.id), impl
  ;;

  let of_async_one_way_rpc ~monitor ~f proto =
    let impl = Function_one_way.make_impl ~monitor ~f proto in
    T (Fn.id, One_way proto, Fn.id), impl
  ;;

  let run_internal
    (type worker query response)
    (t_internal : (worker, query, response) t_internal)
    connection
    ~(arg : query)
    : response Or_error.t Deferred.t
    =
    match t_internal with
    | Plain proto -> Rpc.Rpc.dispatch proto connection arg
    | Piped (proto, Type_equal.T) ->
      let%map result = Rpc.Pipe_rpc.dispatch proto connection arg in
      Or_error.join result |> Or_error.map ~f:(fun (reader, _) -> reader)
    | State (proto, Type_equal.T) ->
      let%map result = Rpc.State_rpc.dispatch proto connection arg in
      Or_error.join result |> Or_error.map ~f:(fun (state, reader, _) -> state, reader)
    | One_way proto -> Rpc.One_way.dispatch proto connection arg |> return
    | Reverse_piped ({ worker_rpc; master_rpc = _; master_in_progress }, Type_equal.T) ->
      let query, updates = arg in
      let key = Function_reverse_piped.Id.create () in
      Hashtbl.add_exn master_in_progress ~key ~data:updates;
      let%map result = Rpc.Rpc.dispatch worker_rpc connection (query, key) in
      Hashtbl.remove master_in_progress key;
      result
    | Directly_piped proto ->
      let arg, f = arg in
      let%map result = Rpc.Pipe_rpc.dispatch_iter proto connection arg ~f in
      Or_error.join result |> Or_error.map ~f:(fun id -> Id_directly_piped.T (proto, id))
  ;;

  let run (T (query_f, t_internal, response_f)) connection ~arg =
    run_internal t_internal connection ~arg:(query_f arg) >>| Or_error.map ~f:response_f
  ;;

  let async_log = T (Fn.id, Piped (Async_log_rpc.rpc, Type_equal.T), Fn.id)
  let close_server = T (Fn.id, One_way Close_server_rpc.rpc, Fn.id)

  module For_internal_testing = struct
    let worker_server_rpc_settings =
      T (Fn.id, Plain Worker_server_rpc_settings_rpc.rpc, Fn.id)
    ;;
  end
end

module Daemonize_args = struct
  type args =
    { umask : int option
    ; redirect_stderr : Fd_redirection.t
    ; redirect_stdout : Fd_redirection.t
    }
  [@@deriving sexp]

  type t =
    [ `Don't_daemonize
    | `Daemonize of args
    ]
  [@@deriving sexp]
end

module type Backend = Backend

module type Worker =
  Worker
    with type ('w, 'q, 'r) _function := ('w, 'q, 'r) Function.t
     and type 'w _id_direct := 'w Function.Direct_pipe.Id.t

module type Functions = Functions

module type Creator =
  Creator
    with type ('w, 'q, 'r) _function := ('w, 'q, 'r) Function.t
     and type ('w, 'q, 'r) _direct := ('w, 'q, 'r) Function.Direct_pipe.t

module type Worker_spec =
  Worker_spec
    with type ('w, 'q, 'r) _function := ('w, 'q, 'r) Function.t
     and type ('w, 'q, 'r) _direct := ('w, 'q, 'r) Function.Direct_pipe.t

module Backend_and_settings = struct
  type t = T : (module Backend with type Settings.t = 'a) * 'a -> t

  let backend (T ((module Backend), _)) = (module Backend : Backend)
end

module Backend : sig
  module Settings : sig
    type t [@@deriving bin_io, sexp]

    val with_backend : t -> (module Backend) -> Backend_and_settings.t
  end

  val set_once_and_get_settings : Backend_and_settings.t -> Settings.t
  val assert_already_initialized_with_same_backend : (module Backend) -> unit

  val serve
    :  ?max_message_size:int
    -> ?buffer_age_limit:Writer.buffer_age_limit
    -> ?handshake_timeout:Time_float.Span.t
    -> ?heartbeat_config:Rpc.Connection.Heartbeat_config.t
    -> implementations:'a Rpc.Implementations.t
    -> initial_connection_state:(Socket.Address.Inet.t -> Rpc.Connection.t -> 'a)
    -> where_to_listen:Tcp.Where_to_listen.inet
    -> Settings.t
    -> (Socket.Address.Inet.t, int) Tcp.Server.t Deferred.t

  val with_client
    :  ?implementations:Rpc.Connection.Client_implementations.t
    -> ?max_message_size:int
    -> ?buffer_age_limit:Writer.buffer_age_limit
    -> ?handshake_timeout:Time_float.Span.t
    -> ?heartbeat_config:Rpc.Connection.Heartbeat_config.t
    -> Settings.t
    -> Socket.Address.Inet.t Tcp.Where_to_connect.t
    -> (Rpc.Connection.t -> 'a Deferred.t)
    -> 'a Or_error.t Deferred.t

  val client
    :  ?implementations:Rpc.Connection.Client_implementations.t
    -> ?max_message_size:int
    -> ?buffer_age_limit:Writer.buffer_age_limit
    -> ?handshake_timeout:Time_float.Span.t
    -> ?heartbeat_config:Rpc.Connection.Heartbeat_config.t
    -> ?description:Info.t
    -> Settings.t
    -> Socket.Address.Inet.t Tcp.Where_to_connect.t
    -> Rpc.Connection.t Or_error.t Deferred.t
end = struct
  module Settings = struct
    type t = Sexp.t [@@deriving bin_io, sexp]

    let with_backend t (module Backend : Backend) =
      let t = [%of_sexp: Backend.Settings.t] t in
      Backend_and_settings.T
        ((module Backend : Backend with type Settings.t = Backend.Settings.t), t)
    ;;
  end

  let backend = Set_once.create ()

  let set_once_and_get_settings
    (Backend_and_settings.T ((module Backend), backend_settings))
    : Settings.t
    =
    Set_once.set_exn backend [%here] (module Backend : Backend);
    [%sexp_of: Backend.Settings.t] backend_settings
  ;;

  let get_backend_and_settings_exn backend_settings =
    let backend = Set_once.get_exn backend [%here] in
    Settings.with_backend backend_settings backend
  ;;

  let assert_already_initialized_with_same_backend (module Other_backend : Backend) =
    let (module Backend : Backend) = Set_once.get_exn backend [%here] in
    let backend = Backend.name in
    let other_backend = Other_backend.name in
    if not (String.equal backend other_backend)
    then
      raise_s
        [%message
          "Rpc_parallel was illegally initialized with two different backends in the \
           same executable"
            ~backend
            ~other_backend]
  ;;

  let serve
    ?max_message_size
    ?buffer_age_limit
    ?handshake_timeout
    ?heartbeat_config
    ~implementations
    ~initial_connection_state
    ~where_to_listen
    backend_settings
    =
    let (Backend_and_settings.T ((module Backend), backend_settings)) =
      get_backend_and_settings_exn backend_settings
    in
    Backend.serve
      ?max_message_size
      ?buffer_age_limit
      ?handshake_timeout
      ?heartbeat_config
      ~implementations
      ~initial_connection_state
      ~where_to_listen
      backend_settings
  ;;

  let with_client
    ?implementations
    ?max_message_size
    ?buffer_age_limit
    ?handshake_timeout
    ?heartbeat_config
    backend_settings
    where_to_connect
    f
    =
    let (Backend_and_settings.T ((module Backend), backend_settings)) =
      get_backend_and_settings_exn backend_settings
    in
    Backend.with_client
      ?implementations
      ?max_message_size
      ?buffer_age_limit
      ?handshake_timeout
      ?heartbeat_config
      backend_settings
      where_to_connect
      f
  ;;

  let client
    ?implementations
    ?max_message_size
    ?buffer_age_limit
    ?handshake_timeout
    ?heartbeat_config
    ?description
    backend_settings
    where_to_connect
    =
    let (Backend_and_settings.T ((module Backend), backend_settings)) =
      get_backend_and_settings_exn backend_settings
    in
    Backend.client
      ?implementations
      ?max_message_size
      ?buffer_age_limit
      ?handshake_timeout
      ?heartbeat_config
      ?description
      backend_settings
      where_to_connect
  ;;
end

module Worker_config = struct
  type t =
    { worker_type : Worker_type_id.t
    ; worker_id : Worker_id.t
    ; name : string option
    ; master : Host_and_port.t
    ; app_rpc_settings : Rpc_settings.t
    ; backend_settings : Backend.Settings.t
    ; cd : string
    ; daemonize_args : Daemonize_args.t
    ; connection_timeout : Time_float.Span.t
    ; worker_command_args : Worker_command_args.t
    }
  [@@deriving fields ~getters, sexp]
end

module Worker_env = struct
  type t =
    { config : Worker_config.t
    ; maybe_release_daemon : unit -> unit
    }
  [@@deriving fields ~getters]
end

(* We want to make sure the [Rpc_settings] used on both the client and server side of all
   communication within a single rpc_parallel application are the same. To help prevent
   mistakes, we write small wrappers around the [Rpc.Connection] functions. *)
let rpc_connection_with_client
  backend_settings
  ~rpc_settings
  ?implementations
  where_to_connect
  f
  =
  let { Rpc_settings.max_message_size
      ; buffer_age_limit
      ; handshake_timeout
      ; heartbeat_config
      }
    =
    rpc_settings
  in
  Backend.with_client
    ?max_message_size
    ?buffer_age_limit
    ?handshake_timeout
    ?heartbeat_config
    ?implementations
    backend_settings
    where_to_connect
    f
;;

let rpc_connection_client backend_settings ~rpc_settings ?implementations where_to_connect
  =
  let { Rpc_settings.max_message_size
      ; buffer_age_limit
      ; handshake_timeout
      ; heartbeat_config
      }
    =
    rpc_settings
  in
  Backend.client
    ?max_message_size
    ?buffer_age_limit
    ?handshake_timeout
    ?heartbeat_config
    ?implementations
    backend_settings
    where_to_connect
;;

let start_server
  backend_settings
  ~rpc_settings
  ~where_to_listen
  ~implementations
  ~initial_connection_state
  =
  let { Rpc_settings.max_message_size
      ; buffer_age_limit
      ; handshake_timeout
      ; heartbeat_config
      }
    =
    rpc_settings
  in
  let implementations =
    Rpc.Implementations.create_exn
      ~implementations:(Versioned_rpc.Menu.add implementations)
      ~on_unknown_rpc:`Close_connection
  in
  let%bind server =
    Backend.serve
      ~implementations
      ~initial_connection_state
      ?max_message_size
      ?buffer_age_limit
      ?handshake_timeout
      ?heartbeat_config
      ~where_to_listen
      backend_settings
  in
  return { Server_with_rpc_settings.server; rpc_settings }
;;

module Heartbeater_master : sig
  type t [@@deriving bin_io]

  val create
    :  host_and_port:Host_and_port.t
    -> rpc_settings:Rpc_settings.t
    -> backend_settings:Backend.Settings.t
    -> t

  val connect_and_shutdown_on_disconnect_exn : t -> [ `Connected ] Deferred.t
end = struct
  type t =
    { host_and_port : Host_and_port.t
    ; rpc_settings : Rpc_settings.t
    ; backend_settings : Backend.Settings.t
    }
  [@@deriving bin_io]

  let create ~host_and_port ~rpc_settings ~backend_settings =
    { host_and_port; rpc_settings; backend_settings }
  ;;

  let connect_and_wait_for_disconnect_exn
    { host_and_port; rpc_settings; backend_settings }
    =
    match%map
      rpc_connection_client
        backend_settings
        ~rpc_settings
        (Tcp.Where_to_connect.of_host_and_port host_and_port)
    with
    | Error e -> raise (Error.to_exn e)
    | Ok conn ->
      `Connected
        (let%map reason = Rpc.Connection.close_reason ~on_close:`finished conn in
         `Disconnected reason)
  ;;

  let connect_and_shutdown_on_disconnect_exn heartbeater =
    let%bind (`Connected wait_for_disconnect) =
      connect_and_wait_for_disconnect_exn heartbeater
    in
    (wait_for_disconnect
     >>> fun (`Disconnected reason) ->
     [%log.global.error
       "Rpc_parallel: Heartbeater with master lost connection... Shutting down."
         (reason : Info.t)];
     Shutdown.shutdown 254);
    return `Connected
  ;;
end

(* All global state that is needed for a process to act as a master *)
type master_state =
  { (* The [Host_and_port.t] corresponding to one's own master Rpc server. *)
    my_server : Host_and_port.t Deferred.t lazy_t
      (* The rpc settings used universally for all rpc connections *)
  ; app_rpc_settings : Rpc_settings.t
  ; backend_settings : Backend.Settings.t
      (* Used to facilitate timeout of connecting to a spawned worker *)
  ; pending : Host_and_port.t Ivar.t Worker_id.Table.t
      (* Arguments used when spawning a new worker. *)
  ; worker_command_args : Worker_command_args.t
      (* Callbacks for spawned worker exceptions along with the monitor that was current
     when [spawn] was called *)
  ; on_failures : ((Error.t -> unit) * Monitor.t) Worker_id.Table.t
  }

(* All global state that is not specific to worker types is collected here *)
type worker_state =
  { (* Currently running worker servers in this process *)
    my_worker_servers : Server_with_rpc_settings.t Worker_id.Table.t
      (* To facilitate process creation cleanup.*)
  ; initialized : [ `Init_started of [ `Initialized ] Or_error.t Deferred.t ] Set_once.t
  }

type global_state =
  { as_master : master_state
  ; as_worker : worker_state
  }

(* Each running instance has the capability to work as a master. This state includes
   information needed to spawn new workers (my_server, my_rpc_settings, pending,
   worker_command_args), information to handle existing spawned workerd (on_failures), and
   information to handle worker servers that are running in process. *)
let global_state : global_state Set_once.t = Set_once.create ()

let get_state_exn () =
  match Set_once.get global_state with
  | None -> failwith "State should have been set already"
  | Some state -> state
;;

let get_master_state_exn () = (get_state_exn ()).as_master
let get_worker_state_exn () = (get_state_exn ()).as_worker

(* Rpcs implemented by master *)
module Register_rpc = struct
  type t = Worker_id.t * Host_and_port.t [@@deriving bin_io]

  type response =
    [ `Shutdown
    | `Registered
    ]
  [@@deriving bin_io]

  let rpc =
    Rpc.Rpc.create
      ~name:"register_worker_rpc"
      ~version:0
      ~bin_query:bin_t
      ~bin_response
      ~include_in_error_count:Only_on_exn
  ;;

  let implementation =
    Rpc.Rpc.implement rpc (fun () (id, worker_hp) ->
      let global_state = get_master_state_exn () in
      match Hashtbl.find global_state.pending id with
      | None ->
        (* We already returned a failure to the [spawn_worker] caller *)
        return `Shutdown
      | Some ivar ->
        Ivar.fill_exn ivar worker_hp;
        return `Registered)
  ;;
end

module Handle_exn_rpc = struct
  type t =
    { id : Worker_id.t
    ; name : string option
    ; error : Error.t
    }
  [@@deriving bin_io]

  let rpc =
    Rpc.Rpc.create
      ~name:"handle_worker_exn_rpc"
      ~version:0
      ~bin_query:bin_t
      ~bin_response:Unit.bin_t
      ~include_in_error_count:Only_on_exn
  ;;

  let implementation =
    Rpc.Rpc.implement rpc (fun () { id; name; error } ->
      let global_state = get_master_state_exn () in
      let on_failure, monitor = Hashtbl.find_exn global_state.on_failures id in
      let name = Option.value ~default:(Worker_id.to_string id) name in
      let error = Error.tag error ~tag:name in
      (* We can't just run [on_failure error] because this will be caught by the Rpc
         monitor for this implementation. *)
      Scheduler.within ~monitor (fun () -> on_failure error);
      return ())
  ;;
end

(* In order to spawn other workers, you must have an rpc server implementing
   [Register_rpc] and [Handle_exn_rpc] *)
let master_implementations =
  [ Register_rpc.implementation; Handle_exn_rpc.implementation ]
;;

(* Setup some global state necessary to act as a master (i.e. spawn workers). This
   includes starting an Rpc server with [master_implementations] *)
let init_master_state backend_and_settings ~rpc_settings ~worker_command_args =
  let backend_settings = Backend.set_once_and_get_settings backend_and_settings in
  match Set_once.get global_state with
  | Some _state -> failwith "Master state must not be set up twice"
  | None ->
    (* Use [size:1] so there is minimal top-level overhead linking with Rpc_parallel *)
    let pending = Worker_id.Table.create ~size:1 () in
    let on_failures = Worker_id.Table.create ~size:1 () in
    let my_worker_servers = Worker_id.Table.create ~size:1 () in
    (* Lazily start our master rpc server *)
    let my_server =
      lazy
        (let%map server =
           start_server
             backend_settings
             ~rpc_settings
             ~where_to_listen:Tcp.Where_to_listen.of_port_chosen_by_os
             ~implementations:master_implementations
             ~initial_connection_state:(fun _ _ -> ())
         in
         Host_and_port.create
           ~host:(Unix.gethostname ())
           ~port:(Tcp.Server.listening_on server.Server_with_rpc_settings.server))
    in
    let as_master =
      { my_server
      ; app_rpc_settings = rpc_settings
      ; backend_settings
      ; pending
      ; worker_command_args
      ; on_failures
      }
    in
    let as_worker = { my_worker_servers; initialized = Set_once.create () } in
    Set_once.set_exn global_state [%here] { as_master; as_worker }
;;

module Make (S : Worker_spec) = struct
  module Id = Utils.Worker_id

  type t =
    { host_and_port : Host_and_port.t
    ; rpc_settings : Rpc_settings.t
    ; backend_settings : Backend.Settings.t
    ; id : Worker_id.t
    ; name : string option
    }
  [@@deriving bin_io, sexp_of]

  type worker = t

  (* Internally we use [Worker_id.t] for all worker ids, but we want to expose an [Id]
     module that is specific to each worker. *)
  let id t = t.id
  let rpc_settings t = t.rpc_settings

  type worker_state =
    { (* A unique identifier for each application of the [Make] functor.
         Because we are running the same executable and this is supposed to run at the
         top level, the master and the workers agree on these ids *)
      type_ : Worker_type_id.t
        (* Persistent states associated with instances of this worker server *)
    ; states : S.Worker_state.t Worker_id.Table.t
        (* To facilitate cleanup in the [Shutdown_on.Connection_closed] case *)
    ; mutable client_has_connected : bool
        (* Build up a list of all implementations for this worker type *)
    ; mutable
        implementations :
        (S.Worker_state.t, S.Connection_state.t) Utils.Internal_connection_state.t
        Rpc.Implementation.t
        list
    ; mutable master_implementations : unit Rpc.Implementation.t list
    }

  let worker_state =
    { type_ = Worker_type_id.create ()
    ; states = Worker_id.Table.create ~size:1 ()
    ; client_has_connected = false
    ; implementations = []
    ; master_implementations = []
    }
  ;;

  (* Schedule all worker implementations in [Monitor.main] so no exceptions are lost.
     Async log automatically throws its exceptions to [Monitor.main] so we can't make
     our own local monitor. We detach [Monitor.main] and send exceptions back to the
     master. *)
  let monitor = Monitor.main

  (* Rpcs implemented by this worker type. The implementations for some must be below
     because User_functions is defined below (by supplying a [Creator] module) *)
  module Init_worker_state_rpc = struct
    module Worker_shutdown_on = struct
      type t =
        | Heartbeater_connection_timeout of Heartbeater_master.t
        | Connection_closed of { connection_timeout : Time_float.Span.t }
        | Called_shutdown_function
      [@@deriving bin_io]
    end

    type query =
      { worker_shutdown_on : Worker_shutdown_on.t
      ; worker : Worker_id.t (* The process that got spawned *)
      ; arg : S.Worker_state.init_arg
      }
    [@@deriving bin_io]

    let rpc =
      Rpc.Rpc.create
        ~name:(sprintf "worker_init_rpc_%s" (Worker_type_id.to_string worker_state.type_))
        ~version:0
        ~bin_query
        ~bin_response:Unit.bin_t
        ~include_in_error_count:Only_on_exn
    ;;
  end

  module Init_connection_state_rpc = struct
    type query =
      { worker_id : Worker_id.t
      ; worker_shutdown_on_disconnect : bool
      ; arg : S.Connection_state.init_arg
      }
    [@@deriving bin_io]

    let rpc =
      Rpc.Rpc.create
        ~name:
          (sprintf
             "set_connection_state_rpc_%s"
             (Worker_type_id.to_string worker_state.type_))
        ~version:0
        ~bin_query
        ~bin_response:Unit.bin_t
        ~include_in_error_count:Only_on_exn
    ;;
  end

  let run_executable how ~env ~worker_command_args ~input =
    Utils.create_worker_env ~extra:env
    |> return
    >>=? fun env ->
    How_to_run.run how ~env ~worker_command_args
    >>|? fun p ->
    Writer.write_sexp (Process.stdin p) input;
    p
  ;;

  module Function_creator = struct
    type nonrec worker = worker
    type connection_state = S.Connection_state.t
    type worker_state = S.Worker_state.t

    let with_add_impl f =
      let func, impl = f () in
      worker_state.implementations <- impl :: worker_state.implementations;
      func
    ;;

    let create_rpc ?name ~f ~bin_input ~bin_output () =
      with_add_impl (fun () ->
        Function.create_rpc ~monitor ~name ~f ~bin_input ~bin_output)
    ;;

    let create_pipe ?name ?client_pushes_back ~f ~bin_input ~bin_output () =
      with_add_impl (fun () ->
        Function.create_pipe ~monitor ~name ~f ~bin_input ~bin_output ~client_pushes_back)
    ;;

    let create_state ?name ?client_pushes_back ~f ~bin_query ~bin_state ~bin_update () =
      with_add_impl (fun () ->
        Function.create_state
          ~monitor
          ~name
          ~f
          ~bin_query
          ~bin_state
          ~bin_update
          ~client_pushes_back)
    ;;

    let create_direct_pipe ?name ?client_pushes_back ~f ~bin_input ~bin_output () =
      with_add_impl (fun () ->
        Function.create_direct_pipe
          ~monitor
          ~name
          ~f
          ~bin_input
          ~bin_output
          ~client_pushes_back)
    ;;

    let create_one_way ?name ~f ~bin_input () =
      with_add_impl (fun () -> Function.create_one_way ~monitor ~name ~f ~bin_input)
    ;;

    let reverse_pipe
      ~create_function
      ?name
      ?client_pushes_back
      ~f
      ~bin_query
      ~bin_update
      ~bin_response
      ()
      =
      let func, `Worker worker_impl, `Master master_impl =
        create_function
          ~monitor
          ~name
          ~f
          ~bin_query
          ~bin_update
          ~bin_response
          ~client_pushes_back
      in
      worker_state.implementations <- worker_impl :: worker_state.implementations;
      worker_state.master_implementations
        <- master_impl :: worker_state.master_implementations;
      func
    ;;

    let create_reverse_pipe
      ?name
      ?client_pushes_back
      ~f
      ~bin_query
      ~bin_update
      ~bin_response
      ()
      =
      reverse_pipe
        ~create_function:Function.create_reverse_pipe
        ?name
        ?client_pushes_back
        ~f
        ~bin_query
        ~bin_update
        ~bin_response
        ()
    ;;

    let create_reverse_direct_pipe
      ?name
      ?client_pushes_back
      ~f
      ~bin_query
      ~bin_update
      ~bin_response
      ()
      =
      reverse_pipe
        ~create_function:Function.create_reverse_direct_pipe
        ?name
        ?client_pushes_back
        ~f
        ~bin_query
        ~bin_update
        ~bin_response
        ()
    ;;

    let of_async_rpc ~f proto =
      with_add_impl (fun () -> Function.of_async_rpc ~monitor ~f proto)
    ;;

    let of_async_pipe_rpc ~f proto =
      with_add_impl (fun () -> Function.of_async_pipe_rpc ~monitor ~f proto)
    ;;

    let of_async_state_rpc ~f proto =
      with_add_impl (fun () -> Function.of_async_state_rpc ~monitor ~f proto)
    ;;

    let of_async_direct_pipe_rpc ~f proto =
      with_add_impl (fun () -> Function.of_async_direct_pipe_rpc ~monitor ~f proto)
    ;;

    let of_async_one_way_rpc ~f proto =
      with_add_impl (fun () -> Function.of_async_one_way_rpc ~monitor ~f proto)
    ;;
  end

  module User_functions = S.Functions (Function_creator)

  let functions = User_functions.functions

  let master_implementations : Rpc.Connection.Client_implementations.t =
    T
      { connection_state = const ()
      ; implementations =
          Rpc.Implementations.create_exn
            ~implementations:(Versioned_rpc.Menu.add worker_state.master_implementations)
            ~on_unknown_rpc:`Close_connection
      }
  ;;

  let serve worker_state_init_arg =
    match Hashtbl.find worker_implementations worker_state.type_ with
    | None ->
      failwith
        "Worker could not find RPC implementations. Make sure the Parallel.Make () \
         functor is applied in the worker. It is suggested to make this toplevel."
    | Some (Worker_implementations.T worker_implementations) ->
      let master_state = get_master_state_exn () in
      let rpc_settings = master_state.app_rpc_settings in
      let backend_settings = master_state.backend_settings in
      let%bind server =
        start_server
          backend_settings
          ~implementations:worker_implementations
          ~initial_connection_state:(fun _address connection ->
            connection, Set_once.create ())
          ~rpc_settings
          ~where_to_listen:Tcp.Where_to_listen.of_port_chosen_by_os
      in
      let id = Worker_id.create () in
      let host = Unix.gethostname () in
      let port = Tcp.Server.listening_on server.server in
      let global_state = get_worker_state_exn () in
      Hashtbl.add_exn global_state.my_worker_servers ~key:id ~data:server;
      let%map state = User_functions.init_worker_state worker_state_init_arg in
      Hashtbl.add_exn worker_state.states ~key:id ~data:state;
      { host_and_port = Host_and_port.create ~host ~port
      ; rpc_settings
      ; backend_settings
      ; id
      ; name = None
      }
  ;;

  module Connection = struct
    type t =
      { connection : Rpc.Connection.t
      ; worker_id : Id.t
      }
    [@@deriving fields ~getters, sexp_of]

    let close t = Rpc.Connection.close t.connection
    let close_finished t = Rpc.Connection.close_finished t.connection
    let close_reason t = Rpc.Connection.close_reason t.connection
    let is_closed t = Rpc.Connection.is_closed t.connection

    let client_aux
      ~worker_shutdown_on_disconnect
      { host_and_port; rpc_settings; id = worker_id; backend_settings; _ }
      init_arg
      =
      match%bind
        rpc_connection_client
          backend_settings
          ~rpc_settings
          ~implementations:master_implementations
          (Tcp.Where_to_connect.of_host_and_port host_and_port)
      with
      | Error error -> return (Error error)
      | Ok connection ->
        (match%bind
           Rpc.Rpc.dispatch
             Init_connection_state_rpc.rpc
             connection
             { worker_id; worker_shutdown_on_disconnect; arg = init_arg }
         with
         | Error e ->
           let%map () = Rpc.Connection.close connection in
           Error e
         | Ok () -> Deferred.Or_error.return { connection; worker_id })
    ;;

    let client = client_aux ~worker_shutdown_on_disconnect:false

    let client_with_worker_shutdown_on_disconnect =
      client_aux ~worker_shutdown_on_disconnect:true
    ;;

    let client_exn worker init_arg = client worker init_arg >>| Or_error.ok_exn

    let with_client worker init_arg ~f =
      client worker init_arg
      >>=? fun conn ->
      let%bind result = Monitor.try_with ~run:`Schedule ~rest:`Log (fun () -> f conn) in
      let%map () = close conn in
      Result.map_error result ~f:(fun exn -> Error.of_exn exn)
    ;;

    let run t ~f ~arg = Function.run f t.connection ~arg
    let run_exn t ~f ~arg = run t ~f ~arg >>| Or_error.ok_exn
    let abort t ~id = Function.Direct_pipe.Id.abort id t.connection
  end

  module Shutdown_on (M : T1) = struct
    type _ t =
      | Connection_closed
          : (connection_state_init_arg:S.Connection_state.init_arg
             -> Connection.t M.t Deferred.t)
            t
      | Heartbeater_connection_timeout : worker M.t Deferred.t t
      | Called_shutdown_function : worker M.t Deferred.t t

    let to_init_worker_state_rpc_query_arg (type a) (t : a t) =
      Staged.stage (fun ~heartbeater_master ~connection_timeout ->
        match t with
        | Heartbeater_connection_timeout ->
          Init_worker_state_rpc.Worker_shutdown_on.Heartbeater_connection_timeout
            heartbeater_master
        | Connection_closed ->
          Init_worker_state_rpc.Worker_shutdown_on.Connection_closed
            { connection_timeout }
        | Called_shutdown_function ->
          Init_worker_state_rpc.Worker_shutdown_on.Called_shutdown_function)
    ;;
  end

  type 'a with_spawn_args =
    ?how:How_to_run.t
    -> ?name:string
    -> ?env:(string * string) list
    -> ?connection_timeout:Time_float.Span.t
    -> ?cd:string
    -> on_failure:(Error.t -> unit)
    -> 'a

  (* This timeout serves three purposes.

     [spawn] returns an error if:
     (1) the master hasn't gotten a register rpc from the spawned worker within
     [connection_timeout] of sending the register rpc.
     (2) a worker hasn't gotten its [init_arg] from the master within [connection_timeout]
     of sending the register rpc

     Additionally, if [~shutdown_on:Connection_closed] was used:
     (3) a worker will shut itself down if it doesn't get a connection from the master
     after spawn succeeded. *)
  let connection_timeout_default = sec 10.

  let decorate_error_if_running_test error =
    let tag =
      "You must call [Rpc_parallel.For_testing.initialize] at the top level before any \
       tests are defined. See lib/rpc_parallel/src/parallel_intf.ml for more \
       information."
    in
    if Core.am_running_test then Error.tag ~tag error else error
  ;;

  (* Specific environment variables that influence process execution that a child should
     inherit from the parent. Not copying can result in confusing behavior based on how
     RPC Parallel is configured because:
     - If no use of other machines is configured, all processes run locally and inherit
       the parent's full environment
     - If another machine is configured, all processes (even local ones) do not inherit
       the parent's environment *)
  let ocaml_env_vars_to_child = [ "OCAMLRUNPARAM"; "ASYNC_CONFIG" ]

  let ocaml_env_from_parent =
    lazy
      (List.filter_map ocaml_env_vars_to_child ~f:(fun var ->
         Option.map (Unix.getenv var) ~f:(fun value -> var, value)))
  ;;

  let spawn_process ~how ~env ~cd ~name ~connection_timeout ~daemonize_args =
    let how = Option.value how ~default:How_to_run.local in
    let env =
      (* [force ocaml_env_from_parent] must come before [Option.value env ~default:[]]
         so any user-supplied [env] can override [ocaml_env_from_parent]. *)
      force ocaml_env_from_parent @ Option.value env ~default:[]
      |> List.fold ~init:String.Map.empty ~f:(fun acc (key, data) ->
           Map.set acc ~key ~data)
      |> Map.to_alist
    in
    let cd = Option.value cd ~default:"/" in
    let connection_timeout =
      Option.value connection_timeout ~default:connection_timeout_default
    in
    (match Set_once.get global_state with
     | None ->
       let error =
         Error.of_string
           "You must initialize this process to run as a master before calling [spawn]. \
            Either use a top-level [start_app] call or use the [Expert] module."
       in
       return (Error (decorate_error_if_running_test error))
     | Some global_state -> Deferred.Or_error.return global_state.as_master)
    >>=? fun global_state ->
    (* generate a unique identifier for this worker *)
    let id = Worker_id.create () in
    let%bind master_server = Lazy.force global_state.my_server in
    let input =
      { Worker_config.worker_type = worker_state.type_
      ; worker_id = id
      ; name
      ; master = master_server
      ; app_rpc_settings = global_state.app_rpc_settings
      ; backend_settings = global_state.backend_settings
      ; cd
      ; daemonize_args
      ; connection_timeout
      ; worker_command_args = global_state.worker_command_args
      }
      |> Worker_config.sexp_of_t
    in
    let pending_ivar = Ivar.create () in
    let worker_command_args =
      match global_state.worker_command_args with
      | Add_master_pid ->
        "RPC_PARALLEL_WORKER"
        :: sprintf !"child-of-%{Pid}@%s" (Unix.getpid ()) (Unix.gethostname ())
        :: Option.to_list name
      | User_supplied { args; pass_name } ->
        if pass_name then args @ Option.to_list name else args
    in
    Hashtbl.add_exn global_state.pending ~key:id ~data:pending_ivar;
    match%map run_executable how ~env ~worker_command_args ~input with
    | Error _ as err ->
      Hashtbl.remove global_state.pending id;
      err
    | Ok process -> Ok (id, process)
  ;;

  let with_client worker ~f =
    let { host_and_port; rpc_settings; backend_settings; _ } = worker in
    rpc_connection_with_client
      backend_settings
      ~rpc_settings
      ~implementations:master_implementations
      (Tcp.Where_to_connect.of_host_and_port host_and_port)
      f
  ;;

  let shutdown worker =
    with_client worker ~f:(fun conn ->
      Rpc.One_way.dispatch Shutdown_rpc.rpc conn () |> return)
    >>| Or_error.join
  ;;

  let with_shutdown_on_error worker ~f =
    match%bind f () with
    | Ok _ as ret -> return ret
    | Error _ as ret ->
      let%bind (_ : unit Or_error.t) = shutdown worker in
      return ret
  ;;

  let wait_for_connection_and_initialize
    ~name
    ~connection_timeout
    ~on_failure
    ~id
    ~worker_shutdown_on
    init_arg
    =
    let connection_timeout =
      Option.value connection_timeout ~default:connection_timeout_default
    in
    let global_state = get_master_state_exn () in
    let pending_ivar = Hashtbl.find_exn global_state.pending id in
    (* Ensure that we got a register from the worker *)
    match%bind Clock.with_timeout connection_timeout (Ivar.read pending_ivar) with
    | `Timeout ->
      Hashtbl.remove global_state.pending id;
      Deferred.Or_error.error_string "Timed out getting connection from process"
    | `Result host_and_port ->
      Hashtbl.remove global_state.pending id;
      let worker =
        { host_and_port
        ; rpc_settings = global_state.app_rpc_settings
        ; backend_settings = global_state.backend_settings
        ; id
        ; name
        }
      in
      let%bind master_server = Lazy.force global_state.my_server in
      Hashtbl.add_exn
        global_state.on_failures
        ~key:worker.id
        ~data:(on_failure, Monitor.current ());
      with_shutdown_on_error worker ~f:(fun () ->
        match%map
          with_client worker ~f:(fun conn ->
            let heartbeater_master =
              Heartbeater_master.create
                ~host_and_port:master_server
                ~rpc_settings:global_state.app_rpc_settings
                ~backend_settings:global_state.backend_settings
            in
            let worker_shutdown_on =
              Staged.unstage worker_shutdown_on ~heartbeater_master ~connection_timeout
            in
            Rpc.Rpc.dispatch
              Init_worker_state_rpc.rpc
              conn
              { worker_shutdown_on; worker = id; arg = init_arg })
        with
        | Error error ->
          Hashtbl.remove global_state.on_failures worker.id;
          Error error
        | Ok (Error e) ->
          Hashtbl.remove global_state.on_failures worker.id;
          Error e
        | Ok (Ok ()) -> Ok worker)
  ;;

  module Spawn_in_foreground_aux_result = struct
    type 'a t =
      ( 'a * Process.t
      , Error.t * [ `Worker_process of Unix.Exit_or_signal.t Deferred.t option ] )
      Result.t
  end

  module Spawn_in_foreground_aux_shutdown_on = Shutdown_on (Spawn_in_foreground_aux_result)

  let finalize_on_error ~finalize f =
    let%bind result = f () in
    match result with
    | Ok x -> return (Ok x)
    | Error e ->
      let finalized = finalize () in
      return (Error (e, finalized))
  ;;

  let spawn_in_foreground_aux
    (type a)
    ?how
    ?name
    ?env
    ?connection_timeout
    ?cd
    ~on_failure
    ~(shutdown_on : a Spawn_in_foreground_aux_shutdown_on.t)
    worker_state_init_arg
    : a
    =
    let open Deferred.Result.Let_syntax in
    let daemonize_args = `Don't_daemonize in
    let worker_shutdown_on =
      Spawn_in_foreground_aux_shutdown_on.to_init_worker_state_rpc_query_arg shutdown_on
    in
    let with_spawned_worker ~f =
      let%bind id, process =
        spawn_process ~how ~env ~cd ~name ~connection_timeout ~daemonize_args
        |> Deferred.Result.map_error ~f:(fun e -> e, `Worker_process None)
      in
      finalize_on_error
        ~finalize:(fun () ->
          let exit_or_signal =
            let open Deferred.Let_syntax in
            let%bind exit_status = Process.wait process in
            let%map () = Writer.close (Process.stdin process)
            and () = Reader.close (Process.stdout process)
            and () = Reader.close (Process.stderr process) in
            exit_status
          in
          `Worker_process (Some exit_or_signal))
        (fun () ->
          let%bind worker =
            wait_for_connection_and_initialize
              ~name
              ~connection_timeout
              ~on_failure
              ~id
              ~worker_shutdown_on
              worker_state_init_arg
          in
          f (worker, process))
    in
    match shutdown_on with
    | Heartbeater_connection_timeout -> with_spawned_worker ~f:return
    | Connection_closed ->
      fun ~connection_state_init_arg ->
        with_spawned_worker ~f:(fun (worker, process) ->
          (* If [Connection_state.init] raises, [client_internal] will close the Rpc
             connection, causing the worker to shutdown. *)
          let%bind conn =
            Connection.client_with_worker_shutdown_on_disconnect
              worker
              connection_state_init_arg
          in
          return (conn, process))
    | Called_shutdown_function -> with_spawned_worker ~f:return
  ;;

  module Spawn_in_foreground_result = struct
    type 'a t = ('a * Process.t) Or_error.t
  end

  module Spawn_in_foreground_shutdown_on = Shutdown_on (Spawn_in_foreground_result)

  let spawn_in_foreground
    (type a)
    ?how
    ?name
    ?env
    ?connection_timeout
    ?cd
    ~on_failure
    ~(shutdown_on : a Spawn_in_foreground_shutdown_on.t)
    worker_state_init_arg
    : a
    =
    let spawn_in_foreground_aux
      (type a)
      ~(shutdown_on : a Spawn_in_foreground_aux_shutdown_on.t)
      : a
      =
      spawn_in_foreground_aux
        ?how
        ?name
        ?env
        ?connection_timeout
        ?cd
        ~on_failure
        ~shutdown_on
        worker_state_init_arg
    in
    match shutdown_on with
    | Heartbeater_connection_timeout ->
      spawn_in_foreground_aux ~shutdown_on:Heartbeater_connection_timeout
      >>| Result.map_error ~f:fst
    | Connection_closed ->
      fun ~connection_state_init_arg ->
        spawn_in_foreground_aux ~shutdown_on:Connection_closed ~connection_state_init_arg
        >>| Result.map_error ~f:fst
    | Called_shutdown_function ->
      spawn_in_foreground_aux ~shutdown_on:Called_shutdown_function
      >>| Result.map_error ~f:fst
  ;;

  module Spawn_in_foreground_exn_result = struct
    type 'a t = 'a * Process.t
  end

  module Spawn_in_foreground_exn_shutdown_on = Shutdown_on (Spawn_in_foreground_exn_result)

  let spawn_in_foreground_exn
    (type a)
    ?how
    ?name
    ?env
    ?connection_timeout
    ?cd
    ~on_failure
    ~(shutdown_on : a Spawn_in_foreground_exn_shutdown_on.t)
    init_arg
    : a
    =
    let open Spawn_in_foreground_exn_shutdown_on in
    match shutdown_on with
    | Connection_closed ->
      fun ~connection_state_init_arg ->
        spawn_in_foreground
          ?how
          ?name
          ?env
          ?connection_timeout
          ?cd
          ~on_failure
          ~shutdown_on:Connection_closed
          init_arg
          ~connection_state_init_arg
        >>| ok_exn
    | Heartbeater_connection_timeout ->
      spawn_in_foreground
        ?how
        ?name
        ?env
        ?connection_timeout
        ?cd
        ~on_failure
        ~shutdown_on:Heartbeater_connection_timeout
        init_arg
      >>| ok_exn
    | Called_shutdown_function ->
      spawn_in_foreground
        ?how
        ?name
        ?env
        ?connection_timeout
        ?cd
        ~on_failure
        ~shutdown_on:Called_shutdown_function
        init_arg
      >>| ok_exn
  ;;

  let wait_for_daemonization_and_collect_stderr name process =
    let%bind () = Writer.close (Process.stdin process) in
    let%bind exit_or_signal = Process.wait process in
    let%bind () = Reader.close (Process.stdout process) in
    let worker_stderr = Reader.lines (Process.stderr process) in
    let%map () =
      Pipe.iter worker_stderr ~f:(fun line ->
        let line' = sprintf "[WORKER %s STDERR]: %s\n" name line in
        Writer.write (Lazy.force Writer.stderr) line' |> return)
    in
    match exit_or_signal with
    | Ok () -> Ok ()
    | Error _ ->
      let error =
        Error.createf
          "Worker process %s"
          (Unix.Exit_or_signal.to_string_hum exit_or_signal)
      in
      Error (decorate_error_if_running_test error)
  ;;

  module Spawn_shutdown_on = Shutdown_on (Or_error)

  let spawn
    (type a)
    ?how
    ?name
    ?env
    ?connection_timeout
    ?cd
    ~on_failure
    ?umask
    ~(shutdown_on : a Spawn_shutdown_on.t)
    ~redirect_stdout
    ~redirect_stderr
    worker_state_init_arg
    : a
    =
    let daemonize_args =
      `Daemonize { Daemonize_args.umask; redirect_stderr; redirect_stdout }
    in
    let worker_shutdown_on =
      Spawn_shutdown_on.to_init_worker_state_rpc_query_arg shutdown_on
    in
    let spawn_worker () =
      match%bind
        spawn_process ~how ~env ~cd ~name ~connection_timeout ~daemonize_args
      with
      | Error e -> return (Error e)
      | Ok (id, process) ->
        let id_or_name = Option.value ~default:(Worker_id.to_string id) name in
        (match%bind wait_for_daemonization_and_collect_stderr id_or_name process with
         | Error e -> return (Error e)
         | Ok () ->
           wait_for_connection_and_initialize
             ~name
             ~connection_timeout
             ~on_failure
             ~id
             ~worker_shutdown_on
             worker_state_init_arg)
    in
    let open Spawn_shutdown_on in
    match shutdown_on with
    | Heartbeater_connection_timeout -> spawn_worker ()
    | Connection_closed ->
      fun ~connection_state_init_arg ->
        spawn_worker ()
        >>=? fun worker ->
        (* If [Connection_state.init] raises, [client_internal] will close the Rpc
           connection, causing the worker to shutdown. *)
        Connection.client_with_worker_shutdown_on_disconnect
          worker
          connection_state_init_arg
    | Called_shutdown_function -> spawn_worker ()
  ;;

  module Spawn_exn_shutdown_on = Shutdown_on (Monad.Ident)

  let spawn_exn
    (type a)
    ?how
    ?name
    ?env
    ?connection_timeout
    ?cd
    ~on_failure
    ?umask
    ~(shutdown_on : a Spawn_exn_shutdown_on.t)
    ~redirect_stdout
    ~redirect_stderr
    init_arg
    : a
    =
    let open Spawn_exn_shutdown_on in
    match shutdown_on with
    | Connection_closed ->
      fun ~connection_state_init_arg ->
        spawn
          ?how
          ?name
          ?env
          ?connection_timeout
          ?cd
          ~on_failure
          ?umask
          ~shutdown_on:Connection_closed
          ~redirect_stdout
          ~redirect_stderr
          init_arg
          ~connection_state_init_arg
        >>| ok_exn
    | Heartbeater_connection_timeout ->
      spawn
        ?how
        ?name
        ?env
        ?connection_timeout
        ?cd
        ~on_failure
        ?umask
        ~shutdown_on:Heartbeater_connection_timeout
        ~redirect_stdout
        ~redirect_stderr
        init_arg
      >>| ok_exn
    | Called_shutdown_function ->
      spawn
        ?how
        ?name
        ?env
        ?connection_timeout
        ?cd
        ~on_failure
        ?umask
        ~shutdown_on:Called_shutdown_function
        ~redirect_stdout
        ~redirect_stderr
        init_arg
      >>| ok_exn
  ;;

  module Deprecated = struct
    let spawn_and_connect
      ?how
      ?name
      ?env
      ?connection_timeout
      ?cd
      ~on_failure
      ?umask
      ~redirect_stdout
      ~redirect_stderr
      ~connection_state_init_arg
      worker_state_init_arg
      =
      spawn
        ?how
        ?name
        ?env
        ?connection_timeout
        ?cd
        ?umask
        ~redirect_stdout
        ~redirect_stderr
        ~on_failure
        ~shutdown_on:Heartbeater_connection_timeout
        worker_state_init_arg
      >>=? fun worker ->
      with_shutdown_on_error worker ~f:(fun () ->
        Connection.client worker connection_state_init_arg)
      >>| Or_error.map ~f:(fun conn -> worker, conn)
    ;;

    let spawn_and_connect_exn
      ?how
      ?name
      ?env
      ?connection_timeout
      ?cd
      ~on_failure
      ?umask
      ~redirect_stdout
      ~redirect_stderr
      ~connection_state_init_arg
      worker_state_init_arg
      =
      spawn_and_connect
        ?how
        ?name
        ?env
        ?connection_timeout
        ?cd
        ~on_failure
        ?umask
        ~redirect_stdout
        ~redirect_stderr
        ~connection_state_init_arg
        worker_state_init_arg
      >>| ok_exn
    ;;
  end

  module For_internal_testing = struct
    module Spawn_in_foreground_result = Spawn_in_foreground_aux_result

    let spawn_in_foreground = spawn_in_foreground_aux
    let master_app_rpc_settings () = (get_master_state_exn ()).app_rpc_settings
  end

  let init_worker_state_impl =
    Rpc.Rpc.implement
      Init_worker_state_rpc.rpc
      (fun
          ( rpc_connection
          , (_ :
              ( Function_creator.worker_state
              , Function_creator.connection_state )
              Utils.Internal_connection_state.t1
              Set_once.t) )
          { Init_worker_state_rpc.worker_shutdown_on; worker; arg }
          ->
      let init_finished =
        Utils.try_within ~monitor (fun () ->
          let%bind on_connection_closed =
            match worker_shutdown_on with
            | Heartbeater_connection_timeout heartbeater_master ->
              let%map `Connected =
                Heartbeater_master.connect_and_shutdown_on_disconnect_exn
                  heartbeater_master
              in
              `Keep_going
            | Connection_closed { connection_timeout = _ } -> return `Shutdown
            | Called_shutdown_function -> return `Keep_going
          in
          let init_finished = User_functions.init_worker_state arg in
          (match on_connection_closed with
           | `Keep_going -> ()
           | `Shutdown ->
             (* The connection that is used to dispatch this rpc is a short-lived
                    connection that is only used to dispatch this rpc. Under normal
                    operation, this connection will be closed by the master after
                    receiving an RPC response. If the connection closes before we finished
                    running [init_worker_state] then we can force ourselves to shutdown
                    early. This is purely a "shutdown early" bit of code. If we remove
                    this code, the worker will still shutdown eventually, but it will be
                    after worker initialization and after [connection_timeout] when the
                    worker realizes that it hasn't received an initial connection from the
                    master. *)
             Rpc.Connection.close_finished rpc_connection
             >>> fun () ->
             if not (Deferred.is_determined init_finished)
             then (
               [%log.global.error
                 "Rpc_parallel: worker connection closed during init... Shutting down"];
               Shutdown.shutdown 254));
          init_finished)
      in
      Set_once.set_exn
        (get_worker_state_exn ()).initialized
        [%here]
        (`Init_started
          (init_finished >>|? fun (_ : Function_creator.worker_state) -> `Initialized));
      match%map init_finished with
      | Error e -> Error.raise e
      | Ok state ->
        (match worker_shutdown_on with
         | Heartbeater_connection_timeout (_ : Heartbeater_master.t)
         | Called_shutdown_function -> ()
         | Connection_closed { connection_timeout } ->
           Clock.after connection_timeout
           >>> fun () ->
           if not worker_state.client_has_connected
           then (
             [%log.global.error_string
               "Rpc_parallel: worker timed out waiting for client connection... Shutting \
                down"];
             Shutdown.shutdown 254));
        Hashtbl.add_exn worker_state.states ~key:worker ~data:state)
  ;;

  let init_connection_state_impl =
    Rpc.Rpc.implement
      Init_connection_state_rpc.rpc
      (fun
          (connection, internal_conn_state)
          { worker_id; worker_shutdown_on_disconnect; arg = init_arg }
          ->
      worker_state.client_has_connected <- true;
      let worker_state = Hashtbl.find_exn worker_state.states worker_id in
      if worker_shutdown_on_disconnect
      then (
        Rpc.Connection.close_reason ~on_close:`finished connection
        >>> fun reason ->
        [%log.global.info
          "Rpc_parallel: initial client connection closed... Shutting down."
            (reason : Info.t)];
        Shutdown.shutdown 0);
      let%map conn_state =
        Utils.try_within_exn ~monitor (fun () ->
          User_functions.init_connection_state ~connection ~worker_state init_arg)
      in
      Set_once.set_exn
        internal_conn_state
        [%here]
        { Utils.Internal_connection_state.worker_id; conn_state; worker_state })
  ;;

  let shutdown_impl =
    Rpc.One_way.implement Shutdown_rpc.rpc (fun _conn_state () ->
      [%log.global.info_string "Rpc_parallel: Got shutdown rpc... Shutting down."];
      Shutdown.shutdown 0)
  ;;

  let close_server_impl =
    Rpc.One_way.implement Close_server_rpc.rpc (fun (_conn, conn_state) () ->
      let { Utils.Internal_connection_state.worker_id; _ } =
        Set_once.get_exn conn_state [%here]
      in
      let global_state = get_worker_state_exn () in
      match Hashtbl.find global_state.my_worker_servers worker_id with
      | None -> ()
      | Some server ->
        Tcp.Server.close server.server
        >>> fun () ->
        Hashtbl.remove global_state.my_worker_servers worker_id;
        Hashtbl.remove worker_state.states worker_id)
  ;;

  let async_log_impl =
    Rpc.Pipe_rpc.implement Async_log_rpc.rpc (fun _conn_state () ->
      let r, w = Pipe.create () in
      let new_output =
        Log.Output.create
          ~flush:(fun () -> Deferred.ignore_m (Pipe.downstream_flushed w))
          (fun msgs ->
            if not (Pipe.is_closed w)
            then Queue.iter msgs ~f:(fun msg -> Pipe.write_without_pushback w msg);
            return ())
      in
      Log.Global.set_output (new_output :: Log.Global.get_output ());
      (* Remove this new output upon the pipe closing. *)
      upon (Pipe.closed w) (fun () ->
        let new_outputs =
          List.filter (Log.Global.get_output ()) ~f:(fun output ->
            not (phys_equal output new_output))
        in
        Log.Global.set_output new_outputs);
      return (Ok r))
  ;;

  let worker_server_rpc_settings_impl =
    Rpc.Rpc.implement Worker_server_rpc_settings_rpc.rpc (fun (_conn, conn_state) () ->
      let { Utils.Internal_connection_state.worker_id; _ } =
        Set_once.get_exn conn_state [%here]
      in
      let global_state = get_worker_state_exn () in
      match Hashtbl.find global_state.my_worker_servers worker_id with
      | None ->
        raise_s
          [%message "Failed to find state associated with worker" (worker_id : Id.t)]
      | Some server -> return server.rpc_settings)
  ;;

  let () =
    worker_state.implementations
      <- [ init_worker_state_impl
         ; init_connection_state_impl
         ; shutdown_impl
         ; close_server_impl
         ; async_log_impl
         ; worker_server_rpc_settings_impl
         ]
         @ worker_state.implementations;
    Hashtbl.add_exn
      worker_implementations
      ~key:worker_state.type_
      ~data:(Worker_implementations.T worker_state.implementations)
  ;;
end

(* Start an Rpc server based on the implementations defined in the [Make] functor
   for this worker type. Return a [Host_and_port.t] describing the server *)
let worker_main backend_settings ~worker_env =
  let { Worker_env.config; maybe_release_daemon } = worker_env in
  let rpc_settings = Worker_config.app_rpc_settings config in
  let id = Worker_config.worker_id config in
  let register my_host_and_port =
    match%map
      rpc_connection_with_client
        backend_settings
        ~rpc_settings
        (Tcp.Where_to_connect.of_host_and_port config.master)
        (fun conn -> Rpc.Rpc.dispatch Register_rpc.rpc conn (id, my_host_and_port))
    with
    | Error error ->
      failwiths ~here:[%here] "Worker failed to register" error [%sexp_of: Error.t]
    | Ok (Error e) ->
      failwiths ~here:[%here] "Worker failed to register" e [%sexp_of: Error.t]
    | Ok (Ok `Shutdown) -> failwith "Got [`Shutdown] on register"
    | Ok (Ok `Registered) -> ()
  in
  (* We want the following two things to occur:

     (1) Catch exceptions in workers and report them back to the master
     (2) Write the exceptions to stderr *)
  let setup_exception_handling () =
    Scheduler.within (fun () ->
      Monitor.detach_and_get_next_error Monitor.main
      >>> fun exn ->
      (* We must be careful that this code here doesn't raise *)
      rpc_connection_with_client
        backend_settings
        ~rpc_settings
        (Tcp.Where_to_connect.of_host_and_port config.master)
        (fun conn ->
        Rpc.Rpc.dispatch
          Handle_exn_rpc.rpc
          conn
          { id; name = config.name; error = Error.of_exn exn })
      >>> fun _ ->
      [%log.global.error_format !"Rpc_parallel: %{Exn} ... Shutting down." exn];
      Shutdown.shutdown 254)
  in
  (* Ensure we do not leak processes. Make sure we have initialized successfully, meaning
     we have heartbeats with the master established if the user wants them. *)
  let setup_cleanup_on_timeout () =
    Clock.after config.connection_timeout
    >>> fun () ->
    match Set_once.get (get_worker_state_exn ()).initialized with
    | None ->
      [%log.global.error_string
        "Rpc_parallel: Timeout getting Init_worker_state rpc from master... Shutting \
         down."];
      Shutdown.shutdown 254
    | Some (`Init_started initialize_result) ->
      initialize_result
      >>> (function
      | Ok `Initialized -> ()
      | Error e -> Error.raise e)
  in
  match Hashtbl.find worker_implementations config.worker_type with
  | None ->
    failwith
      "Worker could not find RPC implementations. Make sure the Parallel.Make () functor \
       is applied in the worker. It is suggested to make this toplevel."
  | Some (Worker_implementations.T worker_implementations) ->
    start_server
      backend_settings
      ~implementations:worker_implementations
      ~initial_connection_state:(fun _address connection ->
        connection, Set_once.create ())
      ~rpc_settings
      ~where_to_listen:Tcp.Where_to_listen.of_port_chosen_by_os
    >>> fun server ->
    let host = Unix.gethostname () in
    let port = Tcp.Server.listening_on server.server in
    let global_state = get_worker_state_exn () in
    Hashtbl.add_exn global_state.my_worker_servers ~key:id ~data:server;
    register (Host_and_port.create ~host ~port)
    >>> fun () ->
    setup_exception_handling ();
    setup_cleanup_on_timeout ();
    (* Daemonize as late as possible but still before running any user code. This lets
       us read any setup errors from stderr *)
    maybe_release_daemon ()
;;

module Expert = struct
  module Worker_env = Worker_env

  let worker_init_before_async_exn () =
    match Utils.whoami () with
    | `Master ->
      failwith
        "[worker_init_before_async_exn] should not be called in a process that was not \
         spawned."
    | `Worker ->
      if Scheduler.is_running ()
      then
        failwith
          "[worker_init_before_async_exn] must be called before the async scheduler has \
           been started.";
      Utils.clear_env ();
      let config =
        try Sexp.input_sexp In_channel.stdin |> Worker_config.t_of_sexp with
        | exn ->
          raise_s
            [%message
              "Unable to read worker config from stdin. Make sure nothing is read from \
               stdin before [worker_init_before_async_exn] is called."
                (exn : Exn.t)]
      in
      let maybe_release_daemon =
        match config.daemonize_args with
        | `Don't_daemonize ->
          Core_unix.chdir config.cd;
          Fn.id
        | `Daemonize { Daemonize_args.umask; redirect_stderr; redirect_stdout } ->
          (* The worker is started via SSH. We want to go to the background so we can close
             the SSH connection, but not until we've connected back to the master via
             Rpc. This allows us to report any initialization errors to the master via the SSH
             connection. *)
          let redirect_stdout = Utils.to_daemon_fd_redirection redirect_stdout in
          let redirect_stderr = Utils.to_daemon_fd_redirection redirect_stderr in
          Staged.unstage
            (Daemon.daemonize_wait
               ~cd:config.cd
               ~redirect_stdout
               ~redirect_stderr
               ?umask
               ())
      in
      { Worker_env.config; maybe_release_daemon }
  ;;

  let start_worker_server_exn backend worker_env =
    let rpc_settings = Worker_env.config worker_env |> Worker_config.app_rpc_settings in
    let backend_settings =
      Worker_env.config worker_env |> Worker_config.backend_settings
    in
    let worker_command_args =
      Worker_env.config worker_env |> Worker_config.worker_command_args
    in
    let backend_and_settings = Backend.Settings.with_backend backend_settings backend in
    init_master_state backend_and_settings ~rpc_settings ~worker_command_args;
    worker_main backend_settings ~worker_env
  ;;

  let start_master_server_exn
    ?rpc_max_message_size
    ?rpc_buffer_age_limit
    ?rpc_handshake_timeout
    ?rpc_heartbeat_config
    ?(pass_name = true)
    backend_and_settings
    ~worker_command_args
    ()
    =
    match Utils.whoami () with
    | `Worker -> failwith "Do not call [init_master_exn] in a spawned worker"
    | `Master ->
      let rpc_settings =
        Rpc_settings.create_with_env_override
          ~max_message_size:rpc_max_message_size
          ~buffer_age_limit:rpc_buffer_age_limit
          ~handshake_timeout:rpc_handshake_timeout
          ~heartbeat_config:rpc_heartbeat_config
      in
      init_master_state
        backend_and_settings
        ~rpc_settings
        ~worker_command_args:(User_supplied { args = worker_command_args; pass_name })
  ;;

  let worker_command backend_and_settings =
    let open Command.Let_syntax in
    Command.async
      ~summary:"internal use only"
      (let%map_open _name = anon (maybe ("NAME" %: string)) in
       let worker_env = worker_init_before_async_exn () in
       fun () ->
         start_worker_server_exn backend_and_settings worker_env;
         Deferred.never ())
      ~behave_nicely_in_pipeline:false
  ;;
end

module State = struct
  type t = [ `started ]

  let get () = Option.map (Set_once.get global_state) ~f:(fun _ -> `started)
end

module For_testing = struct
  let initialize backend_and_settings here =
    if Core.am_running_test
    then (
      match Utils.whoami () with
      | `Master ->
        For_testing_internal.set_initialize_source_code_position here;
        (match State.get () with
         | Some `started ->
           Backend.assert_already_initialized_with_same_backend
             (Backend_and_settings.backend backend_and_settings)
         | None ->
           init_master_state
             backend_and_settings
             ~rpc_settings:Rpc_settings.default
             ~worker_command_args:Add_master_pid)
      | `Worker ->
        if For_testing_internal.worker_should_initialize here
        then (
          let env = Expert.worker_init_before_async_exn () in
          Expert.start_worker_server_exn
            (Backend_and_settings.backend backend_and_settings)
            env;
          never_returns (Scheduler.go ())))
  ;;
end

let start_app
  ?rpc_max_message_size
  ?rpc_buffer_age_limit
  ?rpc_handshake_timeout
  ?rpc_heartbeat_config
  ?when_parsing_succeeds
  ?complete_subcommands
  ?add_validate_parsing_flag
  backend_and_settings
  command
  =
  match Utils.whoami () with
  | `Worker ->
    let worker_env = Expert.worker_init_before_async_exn () in
    Expert.start_worker_server_exn
      (Backend_and_settings.backend backend_and_settings)
      worker_env;
    never_returns (Scheduler.go ())
  | `Master ->
    let rpc_settings =
      Rpc_settings.create_with_env_override
        ~max_message_size:rpc_max_message_size
        ~buffer_age_limit:rpc_buffer_age_limit
        ~handshake_timeout:rpc_handshake_timeout
        ~heartbeat_config:rpc_heartbeat_config
    in
    init_master_state
      backend_and_settings
      ~rpc_settings
      ~worker_command_args:Add_master_pid;
    Command_unix.run
      ?add_validate_parsing_flag
      ?when_parsing_succeeds
      ?complete_subcommands
      command
;;
OCaml

Innovation. Community. Security.