fix: Fix transformers-related hang. (#19165)

marianogappa · web-flow · commit e151df2d2ee0 · 2024-09-16T15:40:18.000+01:00
fixes cloudquery/cloudquery-issues#2436 Several different users have reported that, since transformers have been implemented, certain syncs are hanging forever. The transformations implementation seems to work well in these cases: - A sync completes with no errors - A sync finishes with errors due to the source - A sync finishes with errors due to the transformation However, when a destination has an error, at least sometimes it hangs forever. Each transformer is an implementation of a `grpc.BidiStreamingClient`, which uses these two methods: ```go Send(*Req) error Recv() (*Res, error) ``` Which in turn call these: ```go // SendMsg is generally called by generated code. On error, SendMsg aborts // the stream. If the error was generated by the client, the status is // returned directly; otherwise, io.EOF is returned and the status of // the stream may be discovered using RecvMsg. // // SendMsg blocks until: // - There is sufficient flow control to schedule m with the transport, or // - The stream is done, or // - The stream breaks. // // SendMsg does not wait until the message is received by the server. An // untimely stream closure may result in lost messages. To ensure delivery, // users should ensure the RPC completed successfully using RecvMsg. // // It is safe to have a goroutine calling SendMsg and another goroutine // calling RecvMsg on the same stream at the same time, but it is not safe // to call SendMsg on the same stream in different goroutines. It is also // not safe to call CloseSend concurrently with SendMsg. // // It is not safe to modify the message after calling SendMsg. Tracing // libraries and stats handlers may use the message lazily. SendMsg(m any) error // RecvMsg blocks until it receives a message into m or the stream is // done. It returns io.EOF when the stream completes successfully. On // any other error, the stream is aborted and the error contains the RPC // status. // // It is safe to have a goroutine calling SendMsg and another goroutine // calling RecvMsg on the same stream at the same time, but it is not // safe to call RecvMsg on the same stream in different goroutines. RecvMsg(m any) error ``` Both `SendMsg` and `RecvMsg` can hang forever, and this is what users are reporting. However we may want to mitigate this issue (let's say the flow of data or the orchestration for closing is wrong), it's probably gonna be fragile if we still rely on those methods not hanging. Thus, what I've done is running those two methods (`Recv()` and `Send()`) asynchronously. The Pipeline thread will invoke these, but it will exit if someone closes the pipeline (by using a `select`). I've tested running @disq 's reproduction test 100 times; it no longer hangs nor panics: #19130 (comment) **It correctly fails on a test error on insert configuration** ![Screenshot 2024-09-16 at 12 46 59](https://github.com/user-attachments/assets/6147c351-be56-4dc0-bc71-04d81f12db82) **It correctly succeeds on a success test configuration** ![Screenshot 2024-09-16 at 12 47 29](https://github.com/user-attachments/assets/5e5f0259-a972-49ff-8446-e0ed888de739) **It still doesn't swallow the last batch, as tested by running a stable sync multiple times and seeing the same resource count** ![Screenshot 2024-09-16 at 12 57 17](https://github.com/user-attachments/assets/73f65f27-1896-4333-8ba9-19d8d3ff5833)
diff --git a/cli/internal/transformerpipeline/pipeline.go b/cli/internal/transformerpipeline/pipeline.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"io"
+	"sync/atomic"
+	"time"
 
 	"github.com/cloudquery/plugin-pb-go/pb/plugin/v3"
 	"golang.org/x/sync/errgroup"
@@ -69,59 +71,96 @@ func (lp *TransformerPipeline) Send(data []byte) error {
 		return errors.New("OnOutput must be registered before Send is called, otherwise what do I do with the transformed data?")
 	}
 
-	if lp.clientWrappers[0].isClosed {
-		return errors.New("cannot send data to a closed pipeline")
+	if lp.clientWrappers[0].isClosed.Load() {
+		return nil
 	}
 
-	return lp.clientWrappers[0].client.Send(&plugin.Transform_Request{Record: data})
+	sendCh := make(chan error)
+
+	// Send can block forever (e.g. if grpc buffer is full), so we run it asynchronously
+	// and check if pipeline is closed every second.
+	go func() {
+		err := lp.clientWrappers[0].client.Send(&plugin.Transform_Request{Record: data})
+		sendCh <- err
+	}()
+
+	select {
+	case err := <-sendCh:
+		return err
+	case <-time.After(1 * time.Second): // Check if pipeline is closed every second
+		if lp.clientWrappers[0].isClosed.Load() {
+			return nil
+		}
+	}
+
+	return nil
 }
 
 func (lp *TransformerPipeline) OnOutput(fn func([]byte) error) error {
 	if fn == nil {
 		return errors.New("argument to OnOutput cannot be nil")
 	}
 	lp.clientWrappers[len(lp.clientWrappers)-1].nextSendFn = func(req *plugin.Transform_Request) error {
-		return fn(req.Record)
+		err := fn(req.Record)
+		if err != nil {
+			// Our undocumented convention is that destination errors are unrecoverable. Thus, at this
+			// point we close the pipeline.
+			lp.Close()
+		}
+		return err
 	}
 	return nil
 }
 
 func (lp *TransformerPipeline) Close() {
-	// Closing the pipeline happens on both source as well as destination close.
-	// Not handling this will result in a close of closed channel panic.
-	if lp.clientWrappers[0].isClosed {
-		return
-	}
-
-	// Close the first transformer. The rest will follow gracefully, otherwise records will be lost.
-	lp.clientWrappers[0].client.CloseSend()
-	lp.clientWrappers[0].isClosed = true
+	// Close() can happen in any goroutine, and closing is not thread safe.
+	// Instead of closing, we set a flag that we check on send/recv.
+	lp.clientWrappers[0].isClosed.Store(true)
 }
 
 type clientWrapper struct {
 	i          int
 	client     plugin.Plugin_TransformClient
 	nextSendFn func(*plugin.Transform_Request) error
 	nextClose  func() error
-	isClosed   bool
+	isClosed   atomic.Bool
 }
 
-func (s clientWrapper) startBlocking() error {
+func (s *clientWrapper) startBlocking() error {
 	if s.nextSendFn == nil {
 		return errors.New("nextSendFn is nil")
 	}
-	for {
-		data, err := s.client.Recv()
-		if err == io.EOF {
-			err := s.nextClose()
-			return err
-		}
-		if err != nil {
-			return err
+
+	recvCh := make(chan *plugin.Transform_Request)
+	errCh := make(chan error)
+
+	// Recv can block forever (e.g. if transformer decides to), so
+	// we run it asynchronously and check if pipeline is closed every second.
+	go func() {
+		for {
+			data, err := s.client.Recv()
+			if err != nil {
+				errCh <- err
+			} else {
+				recvCh <- &plugin.Transform_Request{Record: data.Record}
+			}
 		}
-		if err := s.nextSendFn(
-			&plugin.Transform_Request{Record: data.Record},
-		); err != nil {
+	}()
+
+	for {
+		select {
+		case <-time.After(1 * time.Second): // Check if pipeline is closed every second
+			if s.isClosed.Load() {
+				return s.nextClose()
+			}
+		case req := <-recvCh: // Propagate records to next transformer
+			if err := s.nextSendFn(req); err != nil {
+				return err
+			}
+		case err := <-errCh:
+			if err == io.EOF {
+				return s.nextClose()
+			}
 			return err
 		}
 	}
