Skip to content

Commit 54e8a1f

Browse files
Amogh-BharadwajAmogh-Bharadwaj
committed
fix: use os.Pipe and add test
1 parent 133506c commit 54e8a1f

2 files changed

Lines changed: 313 additions & 22 deletions

File tree

‎flow/connectors/postgres/pgdump_schema.go‎

Lines changed: 63 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -43,41 +43,82 @@ func pipeCommand(
4343
appendTLSEnv(ctx, srcCmd, srcConfig)
4444
appendTLSEnv(ctx, psqlCmd, dstConfig)
4545

46-
// pipe source command stdout -> psql stdin
47-
pipe, err := srcCmd.StdoutPipe()
46+
return runPipeline(srcCmd, psqlCmd, srcBinary, "psql")
47+
}
48+
49+
// runPipeline wires srcCmd's stdout into dstCmd's stdin and waits for both.
50+
func runPipeline(srcCmd, dstCmd *exec.Cmd, srcName, dstName string) error {
51+
pr, pw, err := os.Pipe()
4852
if err != nil {
49-
return fmt.Errorf("failed to create %s stdout pipe: %w", srcBinary, err)
53+
return fmt.Errorf("create pipe: %w", err)
5054
}
51-
psqlCmd.Stdin = pipe
55+
srcCmd.Stdout = pw
56+
dstCmd.Stdin = pr
5257

53-
var srcStderr, psqlStderr bytes.Buffer
58+
var srcStderr, dstStderr bytes.Buffer
5459
srcCmd.Stderr = &srcStderr
55-
psqlCmd.Stderr = &psqlStderr
60+
dstCmd.Stderr = &dstStderr
5661

57-
// start psql first so it's ready to read
58-
if err := psqlCmd.Start(); err != nil {
59-
return fmt.Errorf("failed to start psql: %w", err)
62+
// Start dst first so it's ready to read.
63+
if err := dstCmd.Start(); err != nil {
64+
pr.Close()
65+
pw.Close()
66+
return fmt.Errorf("start %s: %w", dstName, err)
6067
}
68+
// dst now owns the read end in its child process.
69+
pr.Close()
6170

62-
// then start source command which writes to the pipe
6371
if err := srcCmd.Start(); err != nil {
64-
// kill psql since source command failed to start
65-
_ = psqlCmd.Process.Kill()
66-
_ = psqlCmd.Wait()
67-
return fmt.Errorf("failed to start %s: %w", srcBinary, err)
72+
pw.Close()
73+
_ = dstCmd.Process.Kill()
74+
_ = dstCmd.Wait()
75+
return fmt.Errorf("start %s: %w", srcName, err)
76+
}
77+
// src now owns the write end in its child process.
78+
pw.Close()
79+
80+
srcDone := make(chan error, 1)
81+
dstDone := make(chan error, 1)
82+
go func() { srcDone <- srcCmd.Wait() }()
83+
go func() { dstDone <- dstCmd.Wait() }()
84+
85+
var (
86+
srcErr, dstErr error
87+
srcKilled, dstKilled bool
88+
)
89+
for range 2 {
90+
select {
91+
case err := <-srcDone:
92+
srcErr = err
93+
if err != nil && dstCmd.ProcessState == nil {
94+
_ = dstCmd.Process.Kill()
95+
dstKilled = true
96+
}
97+
case err := <-dstDone:
98+
dstErr = err
99+
if srcCmd.ProcessState == nil {
100+
// dst exited (success or failure) while src is still running;
101+
// kill src so it doesn't block on a pipe with no reader.
102+
_ = srcCmd.Process.Kill()
103+
srcKilled = true
104+
}
105+
}
68106
}
69107

70-
// wait for source command to finish (closes the pipe, signaling EOF to psql)
71-
srcErr := srcCmd.Wait()
72-
psqlErr := psqlCmd.Wait()
73-
108+
// Report the original cause, not the side we killed in response.
109+
if dstErr != nil && !dstKilled {
110+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", dstName, dstErr, dstStderr.String())
111+
}
112+
if srcErr != nil && !srcKilled {
113+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", srcName, srcErr, srcStderr.String())
114+
}
115+
// Fallback: both sides killed (e.g. ctx cancel) — surface whichever error we have.
74116
if srcErr != nil {
75-
return fmt.Errorf("%s failed: %w\nstderr: %s", srcBinary, srcErr, srcStderr.String())
117+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", srcName, srcErr, srcStderr.String())
76118
}
77-
if psqlErr != nil {
78-
return fmt.Errorf("psql failed: %w\nstderr: %s", psqlErr, psqlStderr.String())
119+
if dstErr != nil {
120+
return fmt.Errorf("%s failed: %w\nstderr:\n%s", dstName, dstErr, dstStderr.String())
79121
}
80-
81122
return nil
82123
}
83124

‎flow/connectors/postgres/pgdump_schema_test.go‎

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,250 @@
1+
package connpostgres
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"errors"
7+
"os/exec"
8+
"runtime"
9+
"strings"
10+
"testing"
11+
"time"
12+
)
13+
14+
// requireUnix skips the test on platforms without the shell utilities used here.
15+
func requireUnix(t *testing.T) {
16+
t.Helper()
17+
if runtime.GOOS == "windows" {
18+
t.Skip("requires unix shell utilities")
19+
}
20+
}
21+
22+
func TestRunPipeline_HappyPath(t *testing.T) {
23+
requireUnix(t)
24+
ctx := t.Context()
25+
26+
src := exec.CommandContext(ctx, "sh", "-c", "printf 'hello world'")
27+
var dstOut bytes.Buffer
28+
dst := exec.CommandContext(ctx, "cat")
29+
dst.Stdout = &dstOut
30+
31+
if err := runPipeline(src, dst, "src", "dst"); err != nil {
32+
t.Fatalf("unexpected error: %v", err)
33+
}
34+
if got := dstOut.String(); got != "hello world" {
35+
t.Fatalf("dst stdout = %q, want %q", got, "hello world")
36+
}
37+
}
38+
39+
func TestRunPipeline_SrcStartFails(t *testing.T) {
40+
ctx := t.Context()
41+
42+
src := exec.CommandContext(ctx, "/nonexistent/peerdb-test-binary")
43+
dst := exec.CommandContext(ctx, "cat")
44+
45+
err := runPipeline(src, dst, "src", "dst")
46+
if err == nil {
47+
t.Fatal("expected error, got nil")
48+
}
49+
if !strings.Contains(err.Error(), "start src") {
50+
t.Fatalf("error %q does not mention src start failure", err)
51+
}
52+
// dst should have been killed and reaped; ProcessState should be set.
53+
if dst.ProcessState == nil {
54+
t.Fatal("dst was not reaped after src start failure")
55+
}
56+
}
57+
58+
func TestRunPipeline_DstStartFails(t *testing.T) {
59+
ctx := t.Context()
60+
61+
src := exec.CommandContext(ctx, "echo", "hi")
62+
dst := exec.CommandContext(ctx, "/nonexistent/peerdb-test-binary")
63+
64+
err := runPipeline(src, dst, "src", "dst")
65+
if err == nil {
66+
t.Fatal("expected error, got nil")
67+
}
68+
if !strings.Contains(err.Error(), "start dst") {
69+
t.Fatalf("error %q does not mention dst start failure", err)
70+
}
71+
// src must not have been started.
72+
if src.ProcessState != nil {
73+
t.Fatal("src should not have been started when dst failed to start")
74+
}
75+
}
76+
77+
func TestRunPipeline_SrcExitsNonZero(t *testing.T) {
78+
requireUnix(t)
79+
ctx := t.Context()
80+
81+
// write some output then exit with error
82+
src := exec.CommandContext(ctx, "sh", "-c", "echo partial; exit 7")
83+
dst := exec.CommandContext(ctx, "cat")
84+
dst.Stdout = &bytes.Buffer{}
85+
86+
err := runPipeline(src, dst, "src", "dst")
87+
if err == nil {
88+
t.Fatal("expected error, got nil")
89+
}
90+
if !strings.Contains(err.Error(), "src failed") {
91+
t.Fatalf("error %q does not mention src failure", err)
92+
}
93+
}
94+
95+
func TestRunPipeline_DstExitsNonZero(t *testing.T) {
96+
requireUnix(t)
97+
ctx := t.Context()
98+
99+
src := exec.CommandContext(ctx, "sh", "-c", "echo hi")
100+
// exit 3 immediately, ignoring stdin
101+
dst := exec.CommandContext(ctx, "sh", "-c", "exit 3")
102+
103+
err := runPipeline(src, dst, "src", "dst")
104+
if err == nil {
105+
t.Fatal("expected error, got nil")
106+
}
107+
// src succeeded so error must be from dst
108+
if !strings.Contains(err.Error(), "dst failed") {
109+
t.Fatalf("error %q does not mention dst failure", err)
110+
}
111+
}
112+
113+
// TestRunPipeline_SrcFailsWhileDstSlow verifies the deadlock-prevention fix:
114+
// if src exits non-zero while dst is still reading slowly, dst is killed so
115+
// runPipeline returns promptly instead of waiting for dst to finish its work.
116+
func TestRunPipeline_SrcFailsWhileDstSlow(t *testing.T) {
117+
requireUnix(t)
118+
ctx := t.Context()
119+
120+
// src writes a small amount (fits in pipe buffer, no blocking) then exits non-zero.
121+
src := exec.CommandContext(ctx, "sh", "-c", "echo hi; exit 9")
122+
// dst is a single process (no shell-spawned children) that doesn't read stdin
123+
// and won't exit on its own. We expect runPipeline to kill it after src fails.
124+
// Note: we deliberately avoid `sh -c "sleep 30; cat"` here -- when sh forks a
125+
// child, that child inherits sh's stderr fd, and Go's exec.Wait blocks
126+
// draining stderr until the inherited fd is closed (i.e. for the full sleep).
127+
// psql doesn't fork children, so this matches real behavior.
128+
dst := exec.CommandContext(ctx, "sleep", "30")
129+
130+
start := time.Now()
131+
done := make(chan error, 1)
132+
go func() { done <- runPipeline(src, dst, "src", "dst") }()
133+
134+
select {
135+
case err := <-done:
136+
if err == nil {
137+
t.Fatal("expected error from src failure")
138+
}
139+
if !strings.Contains(err.Error(), "src failed") {
140+
t.Fatalf("expected src failure, got %v", err)
141+
}
142+
if elapsed := time.Since(start); elapsed > 5*time.Second {
143+
t.Fatalf("runPipeline took %v -- dst was not killed promptly after src failure", elapsed)
144+
}
145+
case <-time.After(10 * time.Second):
146+
t.Fatal("runPipeline hung -- dst was not killed after src failure")
147+
}
148+
}
149+
150+
// TestRunPipeline_DstExitsWhileSrcWriting verifies the inverse: if dst exits
151+
// early while src is producing lots of data, src is killed so it doesn't hang
152+
// forever blocked on a write to a closed pipe (would normally get SIGPIPE,
153+
// but we explicitly kill to be safe / to surface the failure quickly).
154+
func TestRunPipeline_DstExitsWhileSrcWriting(t *testing.T) {
155+
requireUnix(t)
156+
ctx := t.Context()
157+
158+
// src tries to stream a lot of data
159+
src := exec.CommandContext(ctx, "sh", "-c", "yes peerdb | head -c 10000000")
160+
// dst exits immediately without reading
161+
dst := exec.CommandContext(ctx, "sh", "-c", "exit 2")
162+
163+
start := time.Now()
164+
done := make(chan error, 1)
165+
go func() { done <- runPipeline(src, dst, "src", "dst") }()
166+
167+
select {
168+
case err := <-done:
169+
if err == nil {
170+
t.Fatal("expected error from dst failure")
171+
}
172+
// We prefer dst's error since src's failure is just a downstream symptom.
173+
if !strings.Contains(err.Error(), "dst failed") {
174+
t.Fatalf("expected dst failure, got %v", err)
175+
}
176+
if elapsed := time.Since(start); elapsed > 5*time.Second {
177+
t.Fatalf("runPipeline took %v -- src was not killed promptly after dst exit", elapsed)
178+
}
179+
case <-time.After(10 * time.Second):
180+
t.Fatal("runPipeline hung -- src was not killed after dst exited")
181+
}
182+
}
183+
184+
// TestRunPipeline_LargeStream verifies that streaming more than the kernel
185+
// pipe buffer (typically 64KB on Linux) works without deadlock.
186+
func TestRunPipeline_LargeStream(t *testing.T) {
187+
requireUnix(t)
188+
ctx := t.Context()
189+
190+
const size = 2 * 1024 * 1024 // 2 MiB
191+
src := exec.CommandContext(ctx, "sh", "-c", "yes a | head -c "+itoa(size))
192+
var out bytes.Buffer
193+
dst := exec.CommandContext(ctx, "cat")
194+
dst.Stdout = &out
195+
196+
if err := runPipeline(src, dst, "src", "dst"); err != nil {
197+
t.Fatalf("unexpected error: %v", err)
198+
}
199+
if out.Len() != size {
200+
t.Fatalf("dst received %d bytes, want %d", out.Len(), size)
201+
}
202+
}
203+
204+
func TestRunPipeline_ContextCancel(t *testing.T) {
205+
requireUnix(t)
206+
ctx, cancel := context.WithCancel(t.Context())
207+
208+
// long-running src that ignores stdin
209+
src := exec.CommandContext(ctx, "sh", "-c", "sleep 30")
210+
dst := exec.CommandContext(ctx, "cat")
211+
dst.Stdout = &bytes.Buffer{}
212+
213+
done := make(chan error, 1)
214+
go func() { done <- runPipeline(src, dst, "src", "dst") }()
215+
216+
// give them a moment to start
217+
time.Sleep(100 * time.Millisecond)
218+
cancel()
219+
220+
select {
221+
case err := <-done:
222+
if err == nil {
223+
t.Fatal("expected error after context cancel")
224+
}
225+
// CommandContext kills the process; just ensure we got back.
226+
var exitErr *exec.ExitError
227+
if !errors.As(err, &exitErr) && !strings.Contains(err.Error(), "killed") &&
228+
!strings.Contains(err.Error(), "signal") {
229+
// any non-nil error is acceptable here; we're mostly checking we don't hang
230+
t.Logf("got error after cancel: %v", err)
231+
}
232+
case <-time.After(10 * time.Second):
233+
t.Fatal("runPipeline did not return after context cancel")
234+
}
235+
}
236+
237+
// itoa avoids importing strconv just for one call site in tests.
238+
func itoa(n int) string {
239+
if n == 0 {
240+
return "0"
241+
}
242+
var buf [20]byte
243+
i := len(buf)
244+
for n > 0 {
245+
i--
246+
buf[i] = byte('0' + n%10)
247+
n /= 10
248+
}
249+
return string(buf[i:])
250+
}

0 commit comments

Comments
 (0)