Evict closed readers from the cache. (#30133)

lostluck · web-flow · commit d1d34575ac37 · 2024-01-26T11:36:33.000-08:00
Co-authored-by: lostluck &lt;13907733+lostluck@users.noreply.github.com&gt;
diff --git a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go
@@ -425,16 +425,20 @@ func (c *DataChannel) read(ctx context.Context) {
 					continue // we've already closed this cached reader, skip
 				}
 				r.PTransformDone()
-				if r.Closed() {
-					// Clean up local bookkeeping. We'll never see another message
-					// for it again. We have to be careful not to remove the real
-					// one, because readers may be initialized after we've seen
-					// the full stream.
-					delete(cache, id.instID)
-				}
 			}
 			seenLast = seenLast[:0] // reset for re-use
 			c.mu.Unlock()
+			// Scan through the cache and check for any closed readers, and evict them from the cache.
+			// Readers might be closed out of band from the data messages because we received all data
+			// for all transforms in an instruction before the instruction even begun. However, we can't
+			// know this until we received the Control instruction which knows how many transforms for which
+			// we need to receive data. So we check the cache directly every so often and evict closed
+			// readers. We will never recieve data for these instructions again.
+			for instID, r := range cache {
+				if r.Closed() {
+					delete(cache, instID)
+				}
+			}
 		}
 	}
 }

Original file line number	Diff line number	Diff line change
`@@ -425,16 +425,20 @@ func (c *DataChannel) read(ctx context.Context) {`
`425`	`425`	`continue // we've already closed this cached reader, skip`
`426`	`426`	`}`
`427`	`427`	`r.PTransformDone()`
`428`		`- if r.Closed() {`
`429`		`- // Clean up local bookkeeping. We'll never see another message`
`430`		`- // for it again. We have to be careful not to remove the real`
`431`		`- // one, because readers may be initialized after we've seen`
`432`		`- // the full stream.`
`433`		`- delete(cache, id.instID)`
`434`		`- }`
`435`	`428`	`}`
`436`	`429`	`seenLast = seenLast[:0] // reset for re-use`
`437`	`430`	`c.mu.Unlock()`
	`431`	`+ // Scan through the cache and check for any closed readers, and evict them from the cache.`
	`432`	`+ // Readers might be closed out of band from the data messages because we received all data`
	`433`	`+ // for all transforms in an instruction before the instruction even begun. However, we can't`
	`434`	`+ // know this until we received the Control instruction which knows how many transforms for which`
	`435`	`+ // we need to receive data. So we check the cache directly every so often and evict closed`
	`436`	`+ // readers. We will never recieve data for these instructions again.`
	`437`	`+ for instID, r := range cache {`
	`438`	`+ if r.Closed() {`
	`439`	`+ delete(cache, instID)`
	`440`	`+ }`
	`441`	`+ }`
`438`	`442`	`}`
`439`	`443`	`}`
`440`	`444`	`}`