test: deflake directory watcherx

zepatrik · ory-bot · commit ea02ce9d7ae3 · 2026-01-02T11:58:26.000Z
GitOrigin-RevId: 6b1c66f96f46833273490b22179aadcb5c7ce01c
diff --git a/oryx/watcherx/directory.go b/oryx/watcherx/directory.go
@@ -33,97 +33,130 @@ func WatchDirectory(ctx context.Context, dir string, c EventChannel) (Watcher, e
 		return nil, err
 	}
 
-	d := newDispatcher()
-	go streamDirectoryEvents(ctx, w, c, d.trigger, d.done, dir, subDirs)
-	return d, nil
+	dw := &directoryWatcher{
+		dispatcher: newDispatcher(),
+		c:          c,
+		dir:        dir,
+		subDirs:    subDirs,
+		w:          w,
+	}
+	go dw.streamDirectoryEvents(ctx)
+	return dw, nil
+}
+
+type directoryWatcher struct {
+	*dispatcher
+	c       EventChannel
+	dir     string
+	subDirs map[string]struct{}
+	w       *fsnotify.Watcher
 }
 
-func handleEvent(e fsnotify.Event, w *fsnotify.Watcher, c EventChannel, subDirs map[string]struct{}) {
+func (w *directoryWatcher) handleEvent(ctx context.Context, e fsnotify.Event) {
 	if e.Has(fsnotify.Remove) {
-		if _, ok := subDirs[e.Name]; ok {
+		if _, ok := w.subDirs[e.Name]; ok {
 			// we do not want any event on deletion of a directory
-			delete(subDirs, e.Name)
+			delete(w.subDirs, e.Name)
 			return
 		}
-		c <- &RemoveEvent{
+		w.maybeSend(ctx, &RemoveEvent{
 			source: source(e.Name),
-		}
+		})
 		return
 	} else if e.Has(fsnotify.Write | fsnotify.Create) {
 		if stats, err := os.Stat(e.Name); err != nil {
-			c <- &ErrorEvent{
+			w.maybeSend(ctx, &ErrorEvent{
 				error:  errors.WithStack(err),
 				source: source(e.Name),
-			}
+			})
 			return
 		} else if stats.IsDir() {
-			if err := w.Add(e.Name); err != nil {
-				c <- &ErrorEvent{
+			if err := w.w.Add(e.Name); err != nil {
+				w.maybeSend(ctx, &ErrorEvent{
 					error:  errors.WithStack(err),
 					source: source(e.Name),
-				}
+				})
 			}
-			subDirs[e.Name] = struct{}{}
+			w.subDirs[e.Name] = struct{}{}
 			return
 		}
 
 		//#nosec G304 -- false positive
 		data, err := os.ReadFile(e.Name)
 		if err != nil {
-			c <- &ErrorEvent{
+			w.maybeSend(ctx, &ErrorEvent{
 				error:  err,
 				source: source(e.Name),
-			}
+			})
 		} else {
-			c <- &ChangeEvent{
+			w.maybeSend(ctx, &ChangeEvent{
 				data:   data,
 				source: source(e.Name),
-			}
+			})
 		}
 	}
 }
 
-func streamDirectoryEvents(ctx context.Context, w *fsnotify.Watcher, c EventChannel, sendNow <-chan struct{}, sendNowDone chan<- int, dir string, subDirs map[string]struct{}) {
+func (w *directoryWatcher) maybeSend(ctx context.Context, e Event) bool {
+	select {
+	case <-ctx.Done():
+		return false
+	case w.c <- e:
+		return true
+	}
+}
+
+func (w *directoryWatcher) streamDirectoryEvents(ctx context.Context) {
+	defer func() {
+		close(w.done)
+		close(w.c)
+		_ = w.w.Close()
+	}()
 	for {
 		select {
 		case <-ctx.Done():
-			_ = w.Close()
 			return
-		case e := <-w.Events:
-			handleEvent(e, w, c, subDirs)
-		case <-sendNow:
+		case e := <-w.w.Events:
+			w.handleEvent(ctx, e)
+		case <-w.trigger:
 			var eventsSent int
 
-			if err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
+			if err := filepath.Walk(w.dir, func(path string, info os.FileInfo, err error) error {
 				if err != nil {
 					return err
 				}
 				if !info.IsDir() {
 					//#nosec G304 -- false positive
 					data, err := os.ReadFile(path)
 					if err != nil {
-						c <- &ErrorEvent{
+						if !w.maybeSend(ctx, &ErrorEvent{
 							error:  err,
 							source: source(path),
+						}) {
+							return errors.WithStack(context.Canceled)
 						}
 					} else {
-						c <- &ChangeEvent{
+						if !w.maybeSend(ctx, &ChangeEvent{
 							data:   data,
 							source: source(path),
+						}) {
+							return errors.WithStack(context.Canceled)
 						}
 					}
 					eventsSent++
 				}
 				return nil
 			}); err != nil {
-				c <- &ErrorEvent{
+				if !w.maybeSend(ctx, &ErrorEvent{
 					error:  err,
-					source: source(dir),
+					source: source(w.dir),
+				}) {
+					return
 				}
 				eventsSent++
 			}
 
-			sendNowDone <- eventsSent
+			w.done <- eventsSent
 		}
 	}
 }
