ptyfwd: fix draining on exit (#35752)

Fixes a bug introduced by #35663.
Fixes #35746.

src/run/run.c

@@ -1557,7 +1557,12 @@ static int run_context_reconnect(RunContext *c) {
 static void run_context_check_done(RunContext *c) {
         assert(c);
 
-        if (STRPTR_IN_SET(c->active_state, "inactive", "failed") && !c->has_job)
+        bool done = STRPTR_IN_SET(c->active_state, "inactive", "failed") && !c->has_job;
+
+        if (done && c->forward) /* If the service is gone, it's time to drain the output */
+                done = pty_forward_drain(c->forward);
+
+        if (done)
                 (void) sd_event_exit(c->event, EXIT_SUCCESS);
 }
 

src/shared/ptyfwd.c

@@ -81,6 +81,7 @@ struct PTYForward {
         bool read_from_master:1;
 
         bool done:1;
+        bool drain:1;
 
         bool last_char_set:1;
         char last_char;
@@ -240,6 +241,9 @@ static bool drained(PTYForward *f) {
 
         assert(f);
 
+        if (f->done)
+                return true;
+
         if (f->out_buffer_full > 0)
                 return false;
 
@@ -619,10 +623,8 @@ static int do_shovel(PTYForward *f) {
                 f->out_buffer_size = MALLOC_SIZEOF_SAFE(p);
         }
 
-        while ((f->stdin_readable && f->in_buffer_full <= 0) ||
-               (f->master_writable && f->in_buffer_full > 0) ||
-               (f->master_readable && f->out_buffer_full <= 0) ||
-               (f->stdout_writable && f->out_buffer_full > 0)) {
+        for (;;) {
+                bool did_something = false;
 
                 if (f->stdin_readable && f->in_buffer_full < LINE_MAX) {
 
@@ -652,6 +654,8 @@ static int do_shovel(PTYForward *f) {
 
                                 f->in_buffer_full += (size_t) k;
                         }
+
+                        did_something = true;
                 }
 
                 if (f->master_writable && f->in_buffer_full > 0) {
@@ -673,6 +677,8 @@ static int do_shovel(PTYForward *f) {
                                 memmove(f->in_buffer, f->in_buffer + k, f->in_buffer_full - k);
                                 f->in_buffer_full -= k;
                         }
+
+                        did_something = true;
                 }
 
                 if (f->master_readable && f->out_buffer_full < MIN(f->out_buffer_size, (size_t) LINE_MAX)) {
@@ -702,6 +708,8 @@ static int do_shovel(PTYForward *f) {
                                 if (r < 0)
                                         return log_error_errno(r, "Failed to scan for ANSI sequences: %m");
                         }
+
+                        did_something = true;
                 }
 
                 if (f->stdout_writable && f->out_buffer_write_len > 0) {
@@ -738,7 +746,12 @@ static int do_shovel(PTYForward *f) {
                                 f->out_buffer_full -= k;
                                 f->out_buffer_write_len -= k;
                         }
+
+                        did_something = true;
                 }
+
+                if (!did_something)
+                        break;
         }
 
         if (f->stdin_hangup || f->stdout_hangup || f->master_hangup) {
@@ -750,6 +763,11 @@ static int do_shovel(PTYForward *f) {
                         return pty_forward_done(f, 0);
         }
 
+        /* If we were asked to drain, and there's nothing more to handle from the master, then call the callback
+         * too. */
+        if (f->drain && drained(f))
+                return pty_forward_done(f, 0);
+
         return 0;
 }
 
@@ -831,14 +849,8 @@ static int on_exit_event(sd_event_source *e, void *userdata) {
         assert(e);
         assert(e == f->exit_event_source);
 
-        /* Drain the buffer on exit. */
-
-        if (f->done)
-                return 0;
-
-        for (unsigned trial = 0; trial < 1000; trial++) {
-                if (drained(f))
-                        return pty_forward_done(f, 0);
+        if (!pty_forward_drain(f)) {
+                /* If not drained, try to drain the buffer. */
 
                 if (!f->master_hangup)
                         f->master_writable = f->master_readable = true;
@@ -850,12 +862,9 @@ static int on_exit_event(sd_event_source *e, void *userdata) {
                 r = shovel(f);
                 if (r < 0)
                         return r;
-                if (f->done)
-                        return 0;
         }
 
-        /* If we could not drain, then propagate recognizable error code. */
-        return pty_forward_done(f, -ELOOP);
+        return pty_forward_done(f, 0);
 }
 
 int pty_forward_new(
@@ -1077,6 +1086,20 @@ void pty_forward_set_handler(PTYForward *f, PTYForwardHandler cb, void *userdata
         f->userdata = userdata;
 }
 
+bool pty_forward_drain(PTYForward *f) {
+        assert(f);
+
+        /* Starts draining the forwarder. Specifically:
+         *
+         * - Returns true if there are no unprocessed bytes from the pty, false otherwise
+         *
+         * - Makes sure the handler function is called the next time the number of unprocessed bytes hits zero
+         */
+
+        f->drain = true;
+        return drained(f);
+}
+
 int pty_forward_set_priority(PTYForward *f, int64_t priority) {
         int r;
 

src/shared/ptyfwd.h

@@ -33,6 +33,8 @@ bool pty_forward_get_ignore_vhangup(PTYForward *f);
 
 void pty_forward_set_handler(PTYForward *f, PTYForwardHandler handler, void *userdata);
 
+bool pty_forward_drain(PTYForward *f);
+
 int pty_forward_set_priority(PTYForward *f, int64_t priority);
 
 int pty_forward_set_width_height(PTYForward *f, unsigned width, unsigned height);