WIP

src/main/java/com/flowingcode/vaadin/addons/gridexporter/ConcurrentStreamResourceWriter.java

@@ -1,11 +1,13 @@
 package com.flowingcode.vaadin.addons.gridexporter;
 
+import com.vaadin.flow.component.UI;
 import com.vaadin.flow.server.StreamResourceWriter;
 import com.vaadin.flow.server.VaadinSession;
 import java.io.IOException;
 import java.io.InterruptedIOException;
 import java.io.OutputStream;
 import java.nio.channels.InterruptedByTimeoutException;
+import java.util.Optional;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntFunction;
@@ -147,6 +149,22 @@ public float getCost(VaadinSession session) {
     return DEFAULT_COST;
   }
 
+  /**
+   * Returns the UI associated with the current download.
+   * <p>
+   * This method is used to ensure that the UI is still attached to the current session when a
+   * download is initiated. Implementations should return the appropriate UI instance.
+   * </p>
+   *
+   * @return the {@link UI} instance associated with the current download, or {@code null} if no UI
+   *         is available.
+   */
+  protected abstract UI getUI();
+
+  private UI getAttachedUI() {
+    return Optional.ofNullable(getUI()).filter(UI::isAttached).orElse(null);
+  }
+
   /**
    * Callback method that is invoked when a timeout occurs while trying to acquire a permit for
    * starting a download.
@@ -181,15 +199,19 @@ public final void accept(OutputStream stream, VaadinSession session) throws IOEx
     } else {
 
       try {
-
         int permits;
         float cost = getCost(session);
         synchronized (semaphore) {
           permits = costToPermits(cost, semaphore.maxPermits);
         }
 
+        UI ui = getAttachedUI();
+
         if (semaphore.tryAcquire(permits, getTimeout(), TimeUnit.NANOSECONDS)) {
           try {
+            if (ui != null && getAttachedUI()!=ui) {
+              throw new DetachedIOException();
+            }
             delegate.accept(stream, session);
           } finally {
             semaphore.release(permits);

src/main/java/com/flowingcode/vaadin/addons/gridexporter/GridExporter.java

@@ -22,6 +22,7 @@
 import com.flowingcode.vaadin.addons.fontawesome.FontAwesome;
 import com.flowingcode.vaadin.addons.gridhelpers.GridHelper;
 import com.vaadin.flow.component.ComponentUtil;
+import com.vaadin.flow.component.UI;
 import com.vaadin.flow.component.grid.ColumnPathRenderer;
 import com.vaadin.flow.component.grid.Grid;
 import com.vaadin.flow.component.grid.Grid.Column;
@@ -336,6 +337,11 @@ public long getTimeout() {
         return concurrentDownloadTimeoutNanos;
       }
 
+      @Override
+      protected UI getUI() {
+        return grid.getUI().orElse(null);
+      }
+
       @Override
       protected void onTimeout() {
         fireConcurrentDownloadTimeout();

src/test/java/com/flowingcode/vaadin/addons/gridexporter/ConfigurableConcurrentStreamResourceWriter.java

@@ -1,5 +1,6 @@
 package com.flowingcode.vaadin.addons.gridexporter;
 
+import com.vaadin.flow.component.UI;
 import com.vaadin.flow.server.StreamResourceWriter;
 import com.vaadin.flow.server.VaadinSession;
 
@@ -13,6 +14,7 @@ public ConfigurableConcurrentStreamResourceWriter(StreamResourceWriter delegate)
 
   private float cost = GridExporter.DEFAULT_COST;
   private long timeout = 0L;
+  private UI ui;
 
   @Override
   public float getCost(VaadinSession session) {
@@ -32,4 +34,13 @@ public void setTimeout(long timeout) {
     this.timeout = timeout;
   }
 
+  @Override
+  public UI getUI() {
+    return ui;
+  }
+
+  public void setUi(UI ui) {
+    this.ui = ui;
+  }
+
 }