From 1998bce19ffbe17d7e947eeb80886cacc58a740c Mon Sep 17 00:00:00 2001
From: WindSoilder <WindSoilder@outlook.com>
Date: Mon, 10 Oct 2022 20:32:55 +0800
Subject: [PATCH] avoid freeze for table print (#6688)

* avoid freeze for table print

* make failed_with_proper_exit_code work again

* add test case for table

* fix un-used import on windows
---
 crates/nu-cli/src/util.rs                 | 27 +++++++++--------------
 crates/nu-command/tests/commands/table.rs | 25 +++++++++++++++++++++
 crates/nu-protocol/src/pipeline_data.rs   | 15 ++++++++-----
 tests/shell/pipeline/commands/external.rs |  3 +--
 4 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/crates/nu-cli/src/util.rs b/crates/nu-cli/src/util.rs
index 75cdfa3dc9..2ae77dcc71 100644
--- a/crates/nu-cli/src/util.rs
+++ b/crates/nu-cli/src/util.rs
@@ -231,23 +231,18 @@ pub fn eval_source(
     }
 
     match eval_block(engine_state, stack, &block, input, false, false) {
-        Ok(mut pipeline_data) => {
-            if let PipelineData::ExternalStream { exit_code, .. } = &mut pipeline_data {
-                if let Some(exit_code) = exit_code.take().and_then(|it| it.last()) {
-                    stack.add_env_var("LAST_EXIT_CODE".to_string(), exit_code);
-                } else {
-                    set_last_exit_code(stack, 0);
+        Ok(pipeline_data) => {
+            match pipeline_data.print(engine_state, stack, false, false) {
+                Err(err) => {
+                    let working_set = StateWorkingSet::new(engine_state);
+
+                    report_error(&working_set, &err);
+
+                    return false;
+                }
+                Ok(exit_code) => {
+                    set_last_exit_code(stack, exit_code);
                 }
-            } else {
-                set_last_exit_code(stack, 0);
-            }
-
-            if let Err(err) = pipeline_data.print(engine_state, stack, false, false) {
-                let working_set = StateWorkingSet::new(engine_state);
-
-                report_error(&working_set, &err);
-
-                return false;
             }
 
             // reset vt processing, aka ansi because illbehaved externals can break it
diff --git a/crates/nu-command/tests/commands/table.rs b/crates/nu-command/tests/commands/table.rs
index a3468b5070..a458b09323 100644
--- a/crates/nu-command/tests/commands/table.rs
+++ b/crates/nu-command/tests/commands/table.rs
@@ -135,3 +135,28 @@ fn table_expand_flatten_and_deep_1() {
          ╰───┴───┴───┴────────────────────────────────╯"
     );
 }
+
+#[test]
+#[cfg(not(windows))]
+fn external_with_too_much_stdout_should_not_hang_nu() {
+    use nu_test_support::fs::Stub::FileWithContent;
+    use nu_test_support::pipeline;
+    use nu_test_support::playground::Playground;
+    Playground::setup("external with too much stdout", |dirs, sandbox| {
+        let bytes: usize = 81920;
+        let mut large_file_body = String::with_capacity(bytes);
+        for _ in 0..bytes {
+            large_file_body.push_str("a");
+        }
+        sandbox.with_files(vec![FileWithContent("a_large_file.txt", &large_file_body)]);
+
+        let actual = nu!(
+            cwd: dirs.test(), pipeline(
+            r#"
+                cat a_large_file.txt | table
+            "#
+        ));
+
+        assert_eq!(actual.out, large_file_body);
+    })
+}
diff --git a/crates/nu-protocol/src/pipeline_data.rs b/crates/nu-protocol/src/pipeline_data.rs
index 898c735326..e1b91196ba 100644
--- a/crates/nu-protocol/src/pipeline_data.rs
+++ b/crates/nu-protocol/src/pipeline_data.rs
@@ -424,7 +424,7 @@ impl PipelineData {
         stack: &mut Stack,
         no_newline: bool,
         to_stderr: bool,
-    ) -> Result<(), ShellError> {
+    ) -> Result<i64, ShellError> {
         // If the table function is in the declarations, then we can use it
         // to create the table value that will be printed in the terminal
 
@@ -452,10 +452,13 @@ impl PipelineData {
 
             // Make sure everything has finished
             if let Some(exit_code) = exit_code {
-                let _: Vec<_> = exit_code.into_iter().collect();
+                let mut exit_codes: Vec<_> = exit_code.into_iter().collect();
+                if let Some(Value::Int { val, .. }) = exit_codes.pop() {
+                    return Ok(val);
+                }
             }
 
-            return Ok(());
+            return Ok(0);
         }
 
         match engine_state.find_decl("table".as_bytes(), &[]) {
@@ -474,7 +477,7 @@ impl PipelineData {
             }
         };
 
-        Ok(())
+        Ok(0)
     }
 
     fn write_all_and_flush(
@@ -483,7 +486,7 @@ impl PipelineData {
         config: &Config,
         no_newline: bool,
         to_stderr: bool,
-    ) -> Result<(), ShellError> {
+    ) -> Result<i64, ShellError> {
         for item in self {
             let mut out = if let Value::Error { error } = item {
                 let working_set = StateWorkingSet::new(engine_state);
@@ -506,7 +509,7 @@ impl PipelineData {
             }
         }
 
-        Ok(())
+        Ok(0)
     }
 }
 
diff --git a/tests/shell/pipeline/commands/external.rs b/tests/shell/pipeline/commands/external.rs
index df73b96d16..8781b72efd 100644
--- a/tests/shell/pipeline/commands/external.rs
+++ b/tests/shell/pipeline/commands/external.rs
@@ -314,11 +314,10 @@ mod nu_commands {
     }
 
     #[test]
-    #[ignore = "For now we have no way to check LAST_EXIT_CODE in tests, ignore it for now"]
     fn failed_with_proper_exit_code() {
         Playground::setup("external failed", |dirs, _sandbox| {
             let actual = nu!(cwd: dirs.test(), r#"
-            nu -c "cargo build; print $env.LAST_EXIT_CODE"
+            nu -c "cargo build | complete | get exit_code"
             "#);
 
             // cargo for non rust project's exit code is 101.