1- from dataclasses import dataclass
1+ import json
22import os
33import subprocess
44import time
5- from typing import List , Optional
5+ from dataclasses import dataclass
6+ from typing import List , Optional , Callable
67from urllib .request import urlopen
8+
79from . import COMPOSE_FILE
8- import json
10+
911
1012def restart_docker ():
13+ """
14+ Restart all containers defined in the current `COMPOSE_FILE`.
15+
16+ Checks that all spacetimedb containers are up and running after the restart.
17+ If they're not up after a couple of retries, throws an `Exception`.
18+ """
19+ print ("Restarting containers" )
20+
1121 docker = DockerManager (COMPOSE_FILE )
12- # Restart all containers.
1322 docker .compose ("restart" )
14- # Ensure all nodes are reachable from outside.
15- containers = docker .list_containers ()
16- for container in containers :
17- info = json .loads (docker ._execute_command ("docker" , "inspect" , container .name ))
18- try :
19- port = info [0 ]['NetworkSettings' ]['Ports' ]['80/tcp' ][0 ]['HostPort' ]
20- except KeyError :
21- continue
22- ping ("127.0.0.1:{}" .format (port ))
23- # TODO: ping endpoint needs to wait for database startup & leader election
24- time .sleep (2 )
25-
26- def ping (host ):
27- tries = 0
28- while tries < 10 :
29- tries += 1
30- try :
31- print (f"Ping Server at { host } )
32- urlopen (f"http://{ host } )
33- print (f"Server up after { tries } )
34- break
35- except Exception :
36- print ("Server down" )
37- time .sleep (3 )
38- else :
39- raise Exception (f"Server at { host } )
23+ containers = docker .list_spacetimedb_containers ()
24+ if not containers :
25+ raise Exception ("No spacetimedb containers found" )
26+
27+ # Ensure all nodes are running.
28+ attempts = 0
29+ while attempts < 5 :
30+ attempts += 1
31+ if all (container .is_running (docker , spacetimedb_ping_url ) for container in containers ):
32+ # sleep a bit more to allow for leader election etc
33+ # TODO: make ping endpoint consider all server state
34+ time .sleep (2 )
35+ return
36+ else :
37+ time .sleep (1 )
38+
39+ raise Exception ("Not all containers are up and running" )
40+
41+ def spacetimedb_ping_url (port : int ) -> str :
42+ return f"http://127.0.0.1:{ port }
4043
4144@dataclass
4245class DockerContainer :
4346 """Represents a Docker container with its basic properties."""
4447 id : str
4548 name : str
4649
50+ def host_ports (self , docker ) -> set [int ]:
51+ """
52+ Collect all host ports of this container.
53+
54+ Host ports are ports on the host that are bound to ports of the
55+ container.
56+ If the container is not currently running, an empty set is returned.
57+ """
58+ host_ports = set ()
59+ info = docker .inspect_container (self )
60+ for ports in info ['NetworkSettings' ]['Ports' ].values ():
61+ for ip_and_port in ports :
62+ host_port = ip_and_port .get ("HostPort" )
63+ if host_port :
64+ host_ports .add (host_port )
65+ return host_ports
66+
67+ def is_running (self , docker , ping_url : Callable [[int ], str ]) -> bool :
68+ """
69+ Check if the container is running.
70+
71+ `ping_url` takes a port number and returns a URL string that can be used
72+ to determine if the host is running by returning a 200 status.
73+
74+ If `self.host_ports()` returns a non-empty set, and one `ping_url`
75+ request is successful, the container is considered running.
76+ """
77+ host_ports = self .host_ports (docker )
78+ for port in host_ports :
79+ url = ping_url (port )
80+ print (f"Trying { url } , end = '' , flush = True )
81+ try :
82+ with urlopen (url , timeout = 0.2 ) as response :
83+ if response .status == 200 :
84+ print ("ok" )
85+ return True
86+ except Exception as e :
87+ print (f"error: { e } )
88+ continue
89+
90+ print (f"container { self .name } )
91+ return False
92+
4793class DockerManager :
4894 """Manages all Docker and Docker Compose operations."""
4995
@@ -74,19 +120,52 @@ def _execute_command(self, *args: str) -> str:
74120 raise
75121
76122 def compose (self , * args : str ) -> str :
77- """Execute a docker- compose command."""
123+ """Execute a ` docker compose` command."""
78124 return self ._execute_command ("docker" , "compose" , "-f" , self .compose_file , * args )
79125
80- def list_containers (self ) -> List [DockerContainer ]:
81- """List all containers and return as DockerContainer objects."""
82- output = self .compose ("ps" , "-a" , "--format" , "{{.ID}} {{.Name}}" )
126+ def docker (self , * args : str ) -> str :
127+ """Execute a `docker` command."""
128+ return self ._execute_command ("docker" , * args )
129+
130+ def list_containers (self , * filters ) -> List [DockerContainer ]:
131+ """
132+ List the containers of the current compose file and return as DockerContainer objects.
133+
134+ All containers are considered, even if not running ('-a' flag).
135+ The containers may be filtered by 'filters' ('--filter' option).
136+ """
137+ # Use -a so we don't miss a crashed or killed container
138+ # when checking for readiness.
139+ cmd = ["ps" , "-a" ]
140+
141+ # Restrict to the current compose file.
142+ compose_file = os .path .abspath (COMPOSE_FILE )
143+ cmd .extend (["--filter" , f"label=com.docker.compose.project.config_files={ compose_file } ])
144+
145+ # Apply additional filters.
146+ for f in filters :
147+ cmd .extend (["--filter" , f ])
148+
149+ # Output only the fields we need for `DockerContainer`.
150+ cmd .extend (["--format" , "{{.ID}} {{.Names}}" ])
151+
152+ output = self .docker (* cmd )
83153 containers = []
84154 for line in output .splitlines ():
85155 if line .strip ():
86156 container_id , name = line .split (maxsplit = 1 )
87157 containers .append (DockerContainer (id = container_id , name = name ))
88158 return containers
89159
160+ def list_spacetimedb_containers (self ) -> List [DockerContainer ]:
161+ """List all containers running spacetimedb."""
162+ return self .list_containers ("label=app=spacetimedb" )
163+
164+ def inspect_container (self , container : DockerContainer ):
165+ """Run the `inspect` command for `container`, returning the parsed JSON dict."""
166+ info = self .docker ("inspect" , container .name )
167+ return json .loads (info )[0 ]
168+
90169 def get_container_by_name (self , name : str ) -> Optional [DockerContainer ]:
91170 """Find a container by name pattern."""
92171 return next (
@@ -97,29 +176,23 @@ def get_container_by_name(self, name: str) -> Optional[DockerContainer]:
97176 def kill_container (self , container_id : str ):
98177 """Kill a container by ID."""
99178 print (f"Killing container { container_id } )
100- self ._execute_command ( " docker" , "kill" , container_id )
179+ self .docker ( "kill" , container_id )
101180
102181 def start_container (self , container_id : str ):
103182 """Start a container by ID."""
104183 print (f"Starting container { container_id } )
105- self ._execute_command ( " docker" , "start" , container_id )
184+ self .docker ( "start" , container_id )
106185
107186 def disconnect_container (self , container_id : str ):
108187 """Disconnect a container from the network."""
109188 print (f"Disconnecting container { container_id } )
110- self ._execute_command (
111- "docker" , "network" , "disconnect" ,
112- self .network_name , container_id
113- )
189+ self .docker ("network" , "disconnect" , self .network_name , container_id )
114190 print (f"Disconnected container { container_id } )
115191
116192 def connect_container (self , container_id : str ):
117193 """Connect a container to the network."""
118194 print (f"Connecting container { container_id } )
119- self ._execute_command (
120- "docker" , "network" , "connect" ,
121- self .network_name , container_id
122- )
195+ self .docker ("network" , "connect" , self .network_name , container_id )
123196 print (f"Connected container { container_id } )
124197
125198 def generate_root_token (self ) -> str :
0 commit comments