diff --git a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
index a750b01963f6..4eea2c0509dd 100644
--- a/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
+++ b/runners/core-java/src/main/java/org/apache/beam/runners/core/SplittableParDoViaKeyedWorkItems.java
@@ -469,8 +469,9 @@ public String getErrorContext() {
         restrictionState.readLater();
         watermarkEstimatorState.readLater();
         WindowedValue<InputT> read = elementState.read();
+        RestrictionT restriction = restrictionState.read();
         if (timer.causedByDrain() == CausedByDrain.CAUSED_BY_DRAIN) {
-          read =
+          WindowedValue<InputT> drainRead =
               WindowedValues.of(
                   read.getValue(),
                   read.getTimestamp(),
@@ -481,8 +482,73 @@ public String getErrorContext() {
                   CausedByDrain.CAUSED_BY_DRAIN,
                   read.getOpenTelemetryContext(),
                   read.getValueKind());
+          RestrictionTracker.TruncateResult<RestrictionT> truncateResult =
+              invoker.invokeTruncateRestriction(
+                  new BaseArgumentProvider<InputT, OutputT>() {
+                    @Override
+                    public InputT element(DoFn<InputT, OutputT> doFn) {
+                      return drainRead.getValue();
+                    }
+
+                    @Override
+                    public Object restriction() {
+                      return restriction;
+                    }
+
+                    @Override
+                    public RestrictionTracker<?, ?> restrictionTracker() {
+                      return invoker.invokeNewTracker(this);
+                    }
+
+                    @Override
+                    public Instant timestamp(DoFn<InputT, OutputT> doFn) {
+                      return drainRead.getTimestamp();
+                    }
+
+                    @Override
+                    public PipelineOptions pipelineOptions() {
+                      return c.getPipelineOptions();
+                    }
+
+                    @Override
+                    public PaneInfo paneInfo(DoFn<InputT, OutputT> doFn) {
+                      return drainRead.getPaneInfo();
+                    }
+
+                    @Override
+                    public BoundedWindow window() {
+                      return Iterables.getOnlyElement(drainRead.getWindows());
+                    }
+
+                    @Override
+                    public Object sideInput(String tagId) {
+                      PCollectionView<?> view = sideInputMapping.get(tagId);
+                      if (view == null) {
+                        throw new IllegalArgumentException(
+                            "calling getSideInput() with unknown view");
+                      }
+                      return sideInputReader.get(
+                          view, view.getWindowMappingFn().getSideInputWindow(window()));
+                    }
+
+                    @Override
+                    public String getErrorContext() {
+                      return ProcessFn.class.getSimpleName() + ".invokeTruncateRestriction";
+                    }
+                  });
+          if (truncateResult == null) {
+            elementState.clear();
+            restrictionState.clear();
+            watermarkEstimatorState.clear();
+            holdState.clear();
+            return;
+          }
+          RestrictionT truncatedRestriction = truncateResult.getTruncatedRestriction();
+          elementAndRestriction = KV.of(drainRead, truncatedRestriction);
+          restrictionState.write(truncatedRestriction);
+        } else {
+          elementAndRestriction = KV.of(read, restriction);
         }
-        elementAndRestriction = KV.of(read, restrictionState.read());
         watermarkEstimatorStateT = watermarkEstimatorState.read();
       }
 
diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java
index 381e41c98705..247412a85dc4 100644
--- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java
+++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java
@@ -36,6 +36,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.NoSuchElementException;
 import java.util.concurrent.Executors;
 import org.apache.beam.runners.core.SplittableParDoViaKeyedWorkItems.ProcessFn;
@@ -48,12 +49,15 @@
 import org.apache.beam.sdk.state.TimeDomain;
 import org.apache.beam.sdk.testing.ResetDateTimeProvider;
 import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.DoFnTester;
+import org.apache.beam.sdk.transforms.View;
 import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker;
 import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator;
 import org.apache.beam.sdk.transforms.splittabledofn.OffsetRangeTracker;
 import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.IsBounded;
 import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
@@ -152,6 +156,29 @@ private static class ProcessFnTester<
         int maxOutputsPerBundle,
         Duration maxBundleDuration)
         throws Exception {
+      this(
+          currentProcessingTime,
+          fn,
+          inputCoder,
+          restrictionCoder,
+          watermarkEstimatorStateCoder,
+          maxOutputsPerBundle,
+          maxBundleDuration,
+          Collections.emptyMap(),
+          NullSideInputReader.empty());
+    }
+
+    ProcessFnTester(
+        Instant currentProcessingTime,
+        final DoFn<InputT, OutputT> fn,
+        Coder<InputT> inputCoder,
+        Coder<RestrictionT> restrictionCoder,
+        Coder<WatermarkEstimatorStateT> watermarkEstimatorStateCoder,
+        int maxOutputsPerBundle,
+        Duration maxBundleDuration,
+        Map<String, PCollectionView<?>> sideInputMapping,
+        SideInputReader sideInputReader)
+        throws Exception {
       // The exact windowing strategy doesn't matter in this test, but it should be able to
       // encode IntervalWindow's because that's what all tests here use.
       WindowingStrategy<InputT, BoundedWindow> windowingStrategy =
@@ -163,35 +190,20 @@ private static class ProcessFnTester<
               restrictionCoder,
               watermarkEstimatorStateCoder,
               windowingStrategy,
-              Collections.emptyMap());
+              sideInputMapping);
       this.tester = DoFnTester.of(processFn);
       this.timerInternals = new InMemoryTimerInternals();
       this.stateInternals = new TestInMemoryStateInternals<>("dummy");
       processFn.setStateInternalsFactory(key -> stateInternals);
       processFn.setTimerInternalsFactory(key -> timerInternals);
-      processFn.setSideInputReader(NullSideInputReader.empty());
+      processFn.setSideInputReader(sideInputReader);
       processFn.setProcessElementInvoker(
           new OutputAndTimeBoundedSplittableProcessElementInvoker<>(
               fn,
               tester.getPipelineOptions(),
               new DoFnTesterWindowedValueReceiver(tester),
               tester.getMainOutputTag(),
-              new SideInputReader() {
-                @Override
-                public <T> T get(PCollectionView<T> view, BoundedWindow window) {
-                  throw new NoSuchElementException();
-                }
-
-                @Override
-                public <T> boolean contains(PCollectionView<T> view) {
-                  return false;
-                }
-
-                @Override
-                public boolean isEmpty() {
-                  return true;
-                }
-              },
+              sideInputReader,
               Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()),
               maxOutputsPerBundle,
               maxBundleDuration,
@@ -790,4 +802,239 @@ public void testReportsBacklogWithoutGetSize() throws Exception {
       assertEquals(7.0, backlogs.get(0), 0.001);
     }
   }
+
+  private static class TruncateFn extends DoFn<Integer, String> {
+    private final boolean truncateToNull;
+    private final List<String> calls = new ArrayList<>();
+
+    public TruncateFn(boolean truncateToNull) {
+      this.truncateToNull = truncateToNull;
+    }
+
+    @ProcessElement
+    public ProcessContinuation process(
+        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
+      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
+        c.output(c.element() + ":" + i);
+        if (i == 2) {
+          return resume();
+        }
+      }
+      return stop();
+    }
+
+    @GetInitialRestriction
+    public OffsetRange getInitialRestriction() {
+      return new OffsetRange(0, 10);
+    }
+
+    @NewTracker
+    public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
+      return new OffsetRangeTracker(range);
+    }
+
+    @TruncateRestriction
+    public RestrictionTracker.TruncateResult<OffsetRange> truncate(
+        @Restriction OffsetRange restriction, @Element Integer element) {
+      calls.add("truncate:" + element + ":" + restriction);
+      if (truncateToNull) {
+        return null;
+      }
+      // Truncate so that we only process one more element.
+      return RestrictionTracker.TruncateResult.of(
+          new OffsetRange(restriction.getFrom(), restriction.getFrom() + 1));
+    }
+  }
+
+  @Test
+  public void testTruncateRestrictionOnDrain() throws Exception {
+    TruncateFn fn = new TruncateFn(false);
+    Instant base = Instant.now();
+
+    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
+        new ProcessFnTester<>(
+            base,
+            fn,
+            BigEndianIntegerCoder.of(),
+            SerializableCoder.of(OffsetRange.class),
+            VoidCoder.of(),
+            MAX_OUTPUTS_PER_BUNDLE,
+            MAX_BUNDLE_DURATION)) {
+      tester.startElement(42, new OffsetRange(0, 10));
+      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));
+
+      assertTrue(tester.advanceDrain());
+      assertThat(tester.takeOutputElements(), contains("42:3"));
+      assertEquals(Collections.singletonList("truncate:42:[3, 10)"), fn.calls);
+      assertEquals(null, tester.getWatermarkHold());
+    }
+  }
+
+  @Test
+  public void testTruncateRestrictionReturnsNullOnDrain() throws Exception {
+    TruncateFn fn = new TruncateFn(true);
+    Instant base = Instant.now();
+
+    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
+        new ProcessFnTester<>(
+            base,
+            fn,
+            BigEndianIntegerCoder.of(),
+            SerializableCoder.of(OffsetRange.class),
+            VoidCoder.of(),
+            MAX_OUTPUTS_PER_BUNDLE,
+            MAX_BUNDLE_DURATION)) {
+      tester.startElement(42, new OffsetRange(0, 10));
+      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));
+
+      assertTrue(tester.advanceDrain());
+      assertTrue(tester.takeOutputElements().isEmpty());
+      assertEquals(Collections.singletonList("truncate:42:[3, 10)"), fn.calls);
+      assertEquals(null, tester.getWatermarkHold());
+    }
+  }
+
+  private static class TruncateWithSideInputFn extends DoFn<Integer, String> {
+    private final List<String> calls = new ArrayList<>();
+
+    @ProcessElement
+    public ProcessContinuation process(
+        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
+      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
+        c.output(c.element() + ":" + i);
+        if (i == 2) {
+          return resume();
+        }
+      }
+      return stop();
+    }
+
+    @GetInitialRestriction
+    public OffsetRange getInitialRestriction() {
+      return new OffsetRange(0, 10);
+    }
+
+    @NewTracker
+    public OffsetRangeTracker newTracker(@Restriction OffsetRange range) {
+      return new OffsetRangeTracker(range);
+    }
+
+    @TruncateRestriction
+    public RestrictionTracker.TruncateResult<OffsetRange> truncate(
+        @Restriction OffsetRange restriction,
+        @Element Integer element,
+        @SideInput("sideInput") String sideInput) {
+      calls.add("truncate:" + element + ":" + sideInput + ":" + restriction);
+      return RestrictionTracker.TruncateResult.of(
+          new OffsetRange(restriction.getFrom(), restriction.getFrom() + 1));
+    }
+  }
+
+  @Test
+  public void testTruncateRestrictionWithSideInputOnDrain() throws Exception {
+    TruncateWithSideInputFn fn = new TruncateWithSideInputFn();
+    Instant base = Instant.now();
+    PCollectionView<String> view =
+        TestPipeline.create().apply(Create.of("sideValue")).apply(View.asSingleton());
+    Map<String, PCollectionView<?>> sideInputMapping = Collections.singletonMap("sideInput", view);
+    SideInputReader sideInputReader =
+        new SideInputReader() {
+          @Override
+          public <T> T get(PCollectionView<T> v, BoundedWindow window) {
+            if (v.equals(view)) {
+              return (T) "sideValue";
+            }
+            throw new NoSuchElementException();
+          }
+
+          @Override
+          public <T> boolean contains(PCollectionView<T> v) {
+            return v.equals(view);
+          }
+
+          @Override
+          public boolean isEmpty() {
+            return false;
+          }
+        };
+
+    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
+        new ProcessFnTester<>(
+            base,
+            fn,
+            BigEndianIntegerCoder.of(),
+            SerializableCoder.of(OffsetRange.class),
+            VoidCoder.of(),
+            MAX_OUTPUTS_PER_BUNDLE,
+            MAX_BUNDLE_DURATION,
+            sideInputMapping,
+            sideInputReader)) {
+      tester.startElement(42, new OffsetRange(0, 10));
+      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));
+
+      assertTrue(tester.advanceDrain());
+      assertThat(tester.takeOutputElements(), contains("42:3"));
+      assertEquals(Collections.singletonList("truncate:42:sideValue:[3, 10)"), fn.calls);
+    }
+  }
+
+  private static class UnboundedOffsetRangeTracker extends OffsetRangeTracker {
+    public UnboundedOffsetRangeTracker(OffsetRange range) {
+      super(range);
+    }
+
+    @Override
+    public IsBounded isBounded() {
+      return IsBounded.UNBOUNDED;
+    }
+  }
+
+  // Tests that if we don't override TruncateRestriction, the default TruncateRestriction
+  // implementation is used (which for unbounded restrictions stops processing immediately).
+  private static class DefaultTruncateUnboundedFn extends DoFn<Integer, String> {
+    @ProcessElement
+    public ProcessContinuation process(
+        ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) {
+      for (long i = tracker.currentRestriction().getFrom(); tracker.tryClaim(i); ++i) {
+        c.output(c.element() + ":" + i);
+        if (i == 2) {
+          return resume();
+        }
+      }
+      return stop();
+    }
+
+    @GetInitialRestriction
+    public OffsetRange getInitialRestriction() {
+      return new OffsetRange(0, 10);
+    }
+
+    @NewTracker
+    public RestrictionTracker<OffsetRange, Long> newTracker(@Restriction OffsetRange range) {
+      return new UnboundedOffsetRangeTracker(range);
+    }
+  }
+
+  @Test
+  public void testDefaultTruncateRestrictionUnboundedStopsOnDrain() throws Exception {
+    DefaultTruncateUnboundedFn fn = new DefaultTruncateUnboundedFn();
+    Instant base = Instant.now();
+
+    try (ProcessFnTester<Integer, String, OffsetRange, Long, Void> tester =
+        new ProcessFnTester<>(
+            base,
+            fn,
+            BigEndianIntegerCoder.of(),
+            SerializableCoder.of(OffsetRange.class),
+            VoidCoder.of(),
+            MAX_OUTPUTS_PER_BUNDLE,
+            MAX_BUNDLE_DURATION)) {
+      tester.startElement(42, new OffsetRange(0, 10));
+      assertThat(tester.takeOutputElements(), contains("42:0", "42:1", "42:2"));
+
+      assertTrue(tester.advanceDrain());
+      assertTrue(tester.takeOutputElements().isEmpty());
+      assertEquals(null, tester.getWatermarkHold());
+    }
+  }
 }