Stop stale query work promptly on generation cancel

3 files changed, 17 insertions, 4 deletions

diff --git a/internal/io/fs/readfile.go b/internal/io/fs/readfile.go
index 47a999d..0ec2eca 100644
--- a/internal/io/fs/readfile.go
+++ b/internal/io/fs/readfile.go
@@ -13,6 +13,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/mimecast/dtail/internal/ctxutil"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/io/pool"
@@ -206,7 +207,9 @@ func (f *readFile) read(ctx context.Context, fd *os.File, reader *bufio.Reader,
 			if abortReading == status {
 				return err
 			}
-			time.Sleep(time.Millisecond * 100)
+			if !ctxutil.Sleep(ctx, 100*time.Millisecond) {
+				return nil
+			}
 			continue
 		}
 
diff --git a/internal/io/fs/readfile_processor.go b/internal/io/fs/readfile_processor.go
index 8f56bdd..672d6d8 100644
--- a/internal/io/fs/readfile_processor.go
+++ b/internal/io/fs/readfile_processor.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/mimecast/dtail/internal/ctxutil"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/io/pool"
@@ -77,7 +78,9 @@ func (f *readFile) readWithProcessor(ctx context.Context, fd *os.File, reader *b
 			if abortReading == status {
 				return err
 			}
-			time.Sleep(time.Millisecond * 100)
+			if !ctxutil.Sleep(ctx, 100*time.Millisecond) {
+				return nil
+			}
 			continue
 		}
 
diff --git a/internal/io/fs/readfile_processor_optimized.go b/internal/io/fs/readfile_processor_optimized.go
index 2e880e7..1e553ee 100644
--- a/internal/io/fs/readfile_processor_optimized.go
+++ b/internal/io/fs/readfile_processor_optimized.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/mimecast/dtail/internal/ctxutil"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/io/pool"
@@ -346,6 +347,8 @@ func (f *readFile) tailWithProcessorOptimized(ctx context.Context, fd *os.File,
 				return err
 			}
 
+			waitForMoreData := true
+
 			// EOF handling
 			select {
 			case <-ctx.Done():
@@ -354,8 +357,12 @@ func (f *readFile) tailWithProcessorOptimized(ctx context.Context, fd *os.File,
 				if isTruncated, err := f.truncated(fd); isTruncated {
 					return err
 				}
-			case <-time.After(100 * time.Millisecond):
-				// Continue reading after a short delay
+				waitForMoreData = false
+			default:
+			}
+
+			if waitForMoreData && !ctxutil.Sleep(ctx, 100*time.Millisecond) {
+				return nil
 			}
 		}