@@ -56,6 +56,7 @@ class TransientVm:
5656 set_ssh_port : Optional [int ]
5757 vmstate : Optional [store .VmPersistentState ]
5858 state : TransientVmState
59+ copy_out_done : bool
5960
6061 def __init__ (self , config : configuration .RunConfig , vmstore : store .VmStore ) -> None :
6162 self .config = config
@@ -69,6 +70,7 @@ def __init__(self, config: configuration.RunConfig, vmstore: store.VmStore) -> N
6970 self .data_tempfile = tempfile .TemporaryFile ("wb+" , buffering = 0 )
7071 self .set_ssh_port = None
7172 self .vmstate = None
73+ self .copy_out_done = False
7274
7375 def __use_backend_images (self , names : List [str ]) -> List [store .BackendImageInfo ]:
7476 """Ensure the backend images are download for each image spec in 'names'"""
@@ -88,6 +90,16 @@ def __needs_to_copy_in_files_before_running(self) -> bool:
8890 """
8991 return len (self .config .copy_in_before ) > 0
9092
93+ def __qemu_is_running (self ) -> bool :
94+ return (self .qemu_runner and self .qemu_runner .is_running ()) or False
95+
96+ def transfer (self , host_path : str , guest_path : str , copy_from : bool ) -> None :
97+ """ Perform rsync/scp transfer, assumes guest is up """
98+ print ("ALEX: config is " , type (self .config ), repr (self .config ))
99+ use_rsync = configuration .config_wants_rsync_transfer (self .config )
100+ assert self .ssh_config is not None
101+ ssh .transfer (host_path , guest_path , self .ssh_config , copy_from , use_rsync )
102+
91103 def __copy_in_files (self ) -> None :
92104 """Copies the given files or directories (located on the host) into the VM"""
93105 path_mappings = self .config .copy_in_before
@@ -110,19 +122,25 @@ def __copy_in(self, path_mapping: str) -> None:
110122 if not vm_absolute_path .startswith ("/" ):
111123 raise RuntimeError (f"Absolute path for guest required: { vm_absolute_path } )
112124
113- assert isinstance (self .primary_image , store .FrontendImageInfo )
114- assert self .primary_image .backend is not None
115- logging .info (
116- f"Copying from '{ host_path } { self .primary_image .backend .identifier } { vm_absolute_path }
117- )
125+ if not self .__qemu_is_running ():
126+ assert isinstance (self .primary_image , store .FrontendImageInfo )
127+ assert self .primary_image .backend is not None
128+ logging .info (
129+ f"Copying from '{ host_path } { self .primary_image .backend .identifier } { vm_absolute_path }
130+ )
118131
119- with editor .ImageEditor (
120- self .primary_image .path ,
121- self .config .ssh_timeout ,
122- self .config .qmp_timeout ,
123- self .config .rsync ,
124- ) as edit :
125- edit .copy_in (host_path , vm_absolute_path )
132+ with editor .ImageEditor (
133+ self .primary_image .path ,
134+ self .config .ssh_timeout ,
135+ self .config .qmp_timeout ,
136+ self .config .rsync ,
137+ ) as edit :
138+ edit .copy_in (host_path , vm_absolute_path )
139+ else :
140+ logging .info (
141+ f"Copying from '{ host_path } { vm_absolute_path }
142+ )
143+ self .transfer (host_path , vm_absolute_path , copy_from = False )
126144
127145 def __needs_to_copy_out_files_after_running (self ) -> bool :
128146 """Checks if at least one directory on the VM needs to be copied out
@@ -152,19 +170,25 @@ def __copy_out(self, path_mapping: str) -> None:
152170 if not vm_absolute_path .startswith ("/" ):
153171 raise RuntimeError (f"Absolute path for guest required: { vm_absolute_path } )
154172
155- assert isinstance (self .primary_image , store .FrontendImageInfo )
156- assert self .primary_image .backend is not None
157- logging .info (
158- f"Copying from '{ self .primary_image .backend .identifier } { vm_absolute_path } { host_path }
159- )
173+ if not self .__qemu_is_running ():
174+ assert isinstance (self .primary_image , store .FrontendImageInfo )
175+ assert self .primary_image .backend is not None
176+ logging .info (
177+ f"Copying from '{ self .primary_image .backend .identifier } { vm_absolute_path } { host_path }
178+ )
160179
161- with editor .ImageEditor (
162- self .primary_image .path ,
163- self .config .ssh_timeout ,
164- self .config .qmp_timeout ,
165- self .config .rsync ,
166- ) as edit :
167- edit .copy_out (vm_absolute_path , host_path )
180+ with editor .ImageEditor (
181+ self .primary_image .path ,
182+ self .config .ssh_timeout ,
183+ self .config .qmp_timeout ,
184+ self .config .rsync ,
185+ ) as edit :
186+ edit .copy_out (vm_absolute_path , host_path )
187+ else :
188+ logging .info (
189+ f"Copying from '(EXISTING QEMU):{ vm_absolute_path } { host_path }
190+ )
191+ self .transfer (vm_absolute_path , host_path , copy_from = True )
168192
169193 def __qemu_added_args (self ) -> List [str ]:
170194 new_args = ["-name" , self .name ]
@@ -226,6 +250,12 @@ def __prepare_ssh(self) -> None:
226250 extra_options = self .config .ssh_option ,
227251 )
228252
253+ def __ensure_ssh (self ) -> None :
254+ assert self .ssh_config is not None
255+ client = ssh .SshClient (config = self .ssh_config , command = "exit 0" )
256+ conn = client .connect_stdout (timeout = self .config .ssh_timeout )
257+ conn .wait ()
258+
229259 def __connect_ssh (self ) -> int :
230260 assert self .ssh_config is not None
231261 client = ssh .SshClient (config = self .ssh_config , command = self .config .ssh_command )
@@ -281,8 +311,9 @@ def __qemu_sigchld_handler(self, sig: int, _frame: Any) -> None:
281311 def __post_run (self , returncode : int ) -> None :
282312 self .state = TransientVmState .FINISHED
283313
284- if self .__needs_to_copy_out_files_after_running ():
314+ if self .__needs_to_copy_out_files_after_running () and not self . copy_out_done :
285315 self .__copy_out_files ()
316+ self .copy_out_done = True
286317
287318 # If the config name is None, this is a temporary VM,
288319 # so remove any generated frontend images. However, if the
@@ -342,6 +373,13 @@ def run(self) -> None:
342373
343374 def __do_run (self ) -> None :
344375 self .state = TransientVmState .RUNNING
376+ self .copy_out_done = False
377+
378+ # direct copy-in can only be done with SSH console (for the "before" part and only when requested with --direct-copy)
379+ will_direct_copy_in = (
380+ configuration .config_requires_ssh_console (self .config )
381+ and self .config .direct_copy
382+ )
345383
346384 if not self .__is_stateless ():
347385 assert self .vmstate is not None
@@ -357,7 +395,7 @@ def __do_run(self) -> None:
357395 self .config .extra_image
358396 )
359397
360- if self .__needs_to_copy_in_files_before_running ():
398+ if self .__needs_to_copy_in_files_before_running () and not will_direct_copy_in :
361399 self .__copy_in_files ()
362400
363401 print ("Finished preparation. Starting virtual machine" )
@@ -423,6 +461,10 @@ def __do_run(self) -> None:
423461 self .__prepare_proc_data ()
424462
425463 if configuration .config_requires_ssh_console (self .config ):
464+ if self .__needs_to_copy_in_files_before_running () and will_direct_copy_in :
465+ self .__ensure_ssh ()
466+ self .__copy_in_files ()
467+
426468 # Note that we always return the SSH exit code, even if the guest failed to
427469 # shut down. This ensures the shutdown_timeout=0 case is handled as expected.
428470 # (i.e., it returns the SSH code instead of a QEMU error)
@@ -433,6 +475,17 @@ def __do_run(self) -> None:
433475 # SIGCHLD exit.
434476 self .qemu_should_die = True
435477
478+ if self .__needs_to_copy_out_files_after_running () and self .config .direct_copy :
479+ # If the VM was shutdown or is otherwise inaccessible,
480+ # the copy-out operation will also be attempted in __post_run.
481+ try :
482+ self .__copy_out_files ()
483+ self .copy_out_done = True
484+ except utils .TransientProcessError as e :
485+ logging .error (
486+ "copy_out during existing QEMU session failed: {}" .format (e )
487+ )
488+
436489 try :
437490 # Wait a bit for the guest to finish the shutdown and QEMU to exit
438491 self .qemu_runner .shutdown (timeout = self .config .shutdown_timeout )
0 commit comments