materialize/transactor: cancel load before Destroy() if commit fails

go/protocols/materialize/lifecycle.go

@@ -158,10 +158,24 @@ type LoadIterator struct {
 	request     *Request        // Request read into.
 	awaitDoneCh <-chan struct{} // Signaled when last commit acknowledgment has completed.
 	err         error           // Terminal error.
+	cancel      context.CancelFunc
+	ctx         context.Context
 }
 
 // Context returns the Context of this LoadIterator.
-func (it *LoadIterator) Context() context.Context { return it.stream.Context() }
+func (it *LoadIterator) Context() context.Context { 
+	if it.ctx == nil {
+		it.ctx, it.cancel = context.WithCancel(it.stream.Context())
+	}
+
+	return it.ctx
+}
+
+func (it *LoadIterator) Cancel() {
+	if it.ctx != nil {
+		it.cancel()
+	}
+}
 
 // WaitForAcknowledged returns once the prior transaction has been fully acknowledged.
 // Importantly, upon its return a materialization connector is free to issues loads

go/protocols/materialize/transactor.go

@@ -89,7 +89,6 @@ func RunTransactions(
 	opened Response_Opened,
 	transactor Transactor,
 ) (_err error) {
-
 	defer func() {
 		if _err != nil {
 			logrus.WithError(_err).Error("RunTransactions failed")
@@ -236,6 +235,10 @@ func RunTransactions(
 					return fmt.Errorf("commit failed: %w", awaitErr)
 				}
 				awaitDoneCh = nil
+				// Before calling transactor.Destroy, we need to make sure that the load phase
+				// is gracefully cancelled to allow for graceful shutdown of the underlying
+				// connector and to avoid resource leaks from the load phase (e.g. connections to database)
+				loadIt.Cancel()
 			case <-loadDoneCh:
 				if loadErr != nil && loadErr != io.EOF {
 					return fmt.Errorf("transactor.Load: %w", loadErr)