diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
index 4cdf4ae..67fe8b8 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -446,7 +446,7 @@
     def errorprone_version = "2.3.4"
     def google_clients_version = "1.32.1"
     def google_cloud_bigdataoss_version = "2.2.2"
-    def google_cloud_pubsublite_version = "0.13.2"
+    def google_cloud_pubsublite_version = "1.0.4"
     def google_code_gson_version = "2.8.6"
     def google_oauth_clients_version = "1.31.0"
     // Try to keep grpc_version consistent with gRPC version in google_cloud_platform_libraries_bom
diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go
new file mode 100644
index 0000000..5496d8b
--- /dev/null
+++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache.go
@@ -0,0 +1,215 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package statecache implements the state caching feature described by the
+// Beam Fn API
+//
+// The Beam State API and the intended caching behavior are described here:
+// https://docs.google.com/document/d/1BOozW0bzBuz4oHJEuZNDOHdzaV5Y56ix58Ozrqm2jFg/edit#heading=h.7ghoih5aig5m
+package statecache
+
+import (
+	"sync"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
+)
+
+type token string
+
+// SideInputCache stores a cache of reusable inputs for the purposes of
+// eliminating redundant calls to the runner during execution of ParDos
+// using side inputs.
+//
+// A SideInputCache should be initialized when the SDK harness is initialized,
+// creating storage for side input caching. On each ProcessBundleRequest,
+// the cache will process the list of tokens for cacheable side inputs and
+// be queried when side inputs are requested in bundle execution. Once a
+// new bundle request comes in the valid tokens will be updated and the cache
+// will be re-used. In the event that the cache reaches capacity, a random,
+// currently invalid cached object will be evicted.
+type SideInputCache struct {
+	capacity    int
+	mu          sync.Mutex
+	cache       map[token]exec.ReusableInput
+	idsToTokens map[string]token
+	validTokens map[token]int8 // Maps tokens to active bundle counts
+	metrics     CacheMetrics
+}
+
+type CacheMetrics struct {
+	Hits           int64
+	Misses         int64
+	Evictions      int64
+	InUseEvictions int64
+}
+
+// Init makes the cache map and the map of IDs to cache tokens for the
+// SideInputCache. Should only be called once. Returns an error for
+// non-positive capacities.
+func (c *SideInputCache) Init(cap int) error {
+	if cap <= 0 {
+		return errors.Errorf("capacity must be a positive integer, got %v", cap)
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.cache = make(map[token]exec.ReusableInput, cap)
+	c.idsToTokens = make(map[string]token)
+	c.validTokens = make(map[token]int8)
+	c.capacity = cap
+	return nil
+}
+
+// SetValidTokens clears the list of valid tokens then sets new ones, also updating the mapping of
+// transform and side input IDs to cache tokens in the process. Should be called at the start of every
+// new ProcessBundleRequest. If the runner does not support caching, the passed cache token values
+// should be empty and all get/set requests will silently be no-ops.
+func (c *SideInputCache) SetValidTokens(cacheTokens ...fnpb.ProcessBundleRequest_CacheToken) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	for _, tok := range cacheTokens {
+		// User State caching is currently not supported, so these tokens are ignored
+		if tok.GetUserState() != nil {
+			continue
+		}
+		s := tok.GetSideInput()
+		transformID := s.GetTransformId()
+		sideInputID := s.GetSideInputId()
+		t := token(tok.GetToken())
+		c.setValidToken(transformID, sideInputID, t)
+	}
+}
+
+// setValidToken adds a new valid token for a request into the SideInputCache struct
+// by mapping the transform ID and side input ID pairing to the cache token.
+func (c *SideInputCache) setValidToken(transformID, sideInputID string, tok token) {
+	idKey := transformID + sideInputID
+	c.idsToTokens[idKey] = tok
+	count, ok := c.validTokens[tok]
+	if !ok {
+		c.validTokens[tok] = 1
+	} else {
+		c.validTokens[tok] = count + 1
+	}
+}
+
+// CompleteBundle takes the cache tokens passed to set the valid tokens and decrements their
+// usage count for the purposes of maintaining a valid count of whether or not a value is
+// still in use. Should be called once ProcessBundle has completed.
+func (c *SideInputCache) CompleteBundle(cacheTokens ...fnpb.ProcessBundleRequest_CacheToken) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	for _, tok := range cacheTokens {
+		// User State caching is currently not supported, so these tokens are ignored
+		if tok.GetUserState() != nil {
+			continue
+		}
+		t := token(tok.GetToken())
+		c.decrementTokenCount(t)
+	}
+}
+
+// decrementTokenCount decrements the validTokens entry for
+// a given token by 1. Should only be called when completing
+// a bundle.
+func (c *SideInputCache) decrementTokenCount(tok token) {
+	count := c.validTokens[tok]
+	if count == 1 {
+		delete(c.validTokens, tok)
+	} else {
+		c.validTokens[tok] = count - 1
+	}
+}
+
+func (c *SideInputCache) makeAndValidateToken(transformID, sideInputID string) (token, bool) {
+	idKey := transformID + sideInputID
+	// Check if it's a known token
+	tok, ok := c.idsToTokens[idKey]
+	if !ok {
+		return "", false
+	}
+	return tok, c.isValid(tok)
+}
+
+// QueryCache takes a transform ID and side input ID and checking if a corresponding side
+// input has been cached. A query having a bad token (e.g. one that doesn't make a known
+// token or one that makes a known but currently invalid token) is treated the same as a
+// cache miss.
+func (c *SideInputCache) QueryCache(transformID, sideInputID string) exec.ReusableInput {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	tok, ok := c.makeAndValidateToken(transformID, sideInputID)
+	if !ok {
+		return nil
+	}
+	// Check to see if cached
+	input, ok := c.cache[tok]
+	if !ok {
+		c.metrics.Misses++
+		return nil
+	}
+
+	c.metrics.Hits++
+	return input
+}
+
+// SetCache allows a user to place a ReusableInput materialized from the reader into the SideInputCache
+// with its corresponding transform ID and side input ID. If the IDs do not pair with a known, valid token
+// then we silently do not cache the input, as this is an indication that the runner is treating that input
+// as uncacheable.
+func (c *SideInputCache) SetCache(transformID, sideInputID string, input exec.ReusableInput) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	tok, ok := c.makeAndValidateToken(transformID, sideInputID)
+	if !ok {
+		return
+	}
+	if len(c.cache) >= c.capacity {
+		c.evictElement()
+	}
+	c.cache[tok] = input
+}
+
+func (c *SideInputCache) isValid(tok token) bool {
+	count, ok := c.validTokens[tok]
+	// If the token is not known or not in use, return false
+	return ok && count > 0
+}
+
+// evictElement randomly evicts a ReusableInput that is not currently valid from the cache.
+// It should only be called by a goroutine that obtained the lock in SetCache.
+func (c *SideInputCache) evictElement() {
+	deleted := false
+	// Select a key from the cache at random
+	for k := range c.cache {
+		// Do not evict an element if it's currently valid
+		if !c.isValid(k) {
+			delete(c.cache, k)
+			c.metrics.Evictions++
+			deleted = true
+			break
+		}
+	}
+	// Nothing is deleted if every side input is still valid. Clear
+	// out a random entry and record the in-use eviction
+	if !deleted {
+		for k := range c.cache {
+			delete(c.cache, k)
+			c.metrics.InUseEvictions++
+			break
+		}
+	}
+}
diff --git a/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go
new file mode 100644
index 0000000..b9970c3
--- /dev/null
+++ b/sdks/go/pkg/beam/core/runtime/harness/statecache/statecache_test.go
@@ -0,0 +1,290 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package statecache
+
+import (
+	"testing"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec"
+	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
+)
+
+// TestReusableInput implements the ReusableInput interface for the purposes
+// of testing.
+type TestReusableInput struct {
+	transformID string
+	sideInputID string
+	value       interface{}
+}
+
+func makeTestReusableInput(transformID, sideInputID string, value interface{}) exec.ReusableInput {
+	return &TestReusableInput{transformID: transformID, sideInputID: sideInputID, value: value}
+}
+
+// Init is a ReusableInput interface method, this is a no-op.
+func (r *TestReusableInput) Init() error {
+	return nil
+}
+
+// Value returns the stored value in the TestReusableInput.
+func (r *TestReusableInput) Value() interface{} {
+	return r.value
+}
+
+// Reset clears the value in the TestReusableInput.
+func (r *TestReusableInput) Reset() error {
+	r.value = nil
+	return nil
+}
+
+func TestInit(t *testing.T) {
+	var s SideInputCache
+	err := s.Init(5)
+	if err != nil {
+		t.Errorf("SideInputCache failed but should have succeeded, got %v", err)
+	}
+}
+
+func TestInit_Bad(t *testing.T) {
+	var s SideInputCache
+	err := s.Init(0)
+	if err == nil {
+		t.Error("SideInputCache init succeeded but should have failed")
+	}
+}
+
+func TestQueryCache_EmptyCase(t *testing.T) {
+	var s SideInputCache
+	err := s.Init(1)
+	if err != nil {
+		t.Fatalf("cache init failed, got %v", err)
+	}
+	output := s.QueryCache("side1", "transform1")
+	if output != nil {
+		t.Errorf("Cache hit when it should have missed, got %v", output)
+	}
+}
+
+func TestSetCache_UncacheableCase(t *testing.T) {
+	var s SideInputCache
+	err := s.Init(1)
+	if err != nil {
+		t.Fatalf("cache init failed, got %v", err)
+	}
+	input := makeTestReusableInput("t1", "s1", 10)
+	s.SetCache("t1", "s1", input)
+	output := s.QueryCache("t1", "s1")
+	if output != nil {
+		t.Errorf("Cache hit when should have missed, got %v", output)
+	}
+}
+
+func TestSetCache_CacheableCase(t *testing.T) {
+	var s SideInputCache
+	err := s.Init(1)
+	if err != nil {
+		t.Fatalf("cache init failed, got %v", err)
+	}
+	transID := "t1"
+	sideID := "s1"
+	tok := token("tok1")
+	s.setValidToken(transID, sideID, tok)
+	input := makeTestReusableInput(transID, sideID, 10)
+	s.SetCache(transID, sideID, input)
+	output := s.QueryCache(transID, sideID)
+	if output == nil {
+		t.Fatalf("call to query cache missed when should have hit")
+	}
+	val, ok := output.Value().(int)
+	if !ok {
+		t.Errorf("failed to convert value to integer, got %v", output.Value())
+	}
+	if val != 10 {
+		t.Errorf("element mismatch, expected 10, got %v", val)
+	}
+}
+
+func makeRequest(transformID, sideInputID string, t token) fnpb.ProcessBundleRequest_CacheToken {
+	var tok fnpb.ProcessBundleRequest_CacheToken
+	var wrap fnpb.ProcessBundleRequest_CacheToken_SideInput_
+	var side fnpb.ProcessBundleRequest_CacheToken_SideInput
+	side.TransformId = transformID
+	side.SideInputId = sideInputID
+	wrap.SideInput = &side
+	tok.Type = &wrap
+	tok.Token = []byte(t)
+	return tok
+}
+
+func TestSetValidTokens(t *testing.T) {
+	inputs := []struct {
+		transformID string
+		sideInputID string
+		tok         token
+	}{
+		{
+			"t1",
+			"s1",
+			"tok1",
+		},
+		{
+			"t2",
+			"s2",
+			"tok2",
+		},
+		{
+			"t3",
+			"s3",
+			"tok3",
+		},
+	}
+
+	var s SideInputCache
+	err := s.Init(3)
+	if err != nil {
+		t.Fatalf("cache init failed, got %v", err)
+	}
+
+	var tokens []fnpb.ProcessBundleRequest_CacheToken
+	for _, input := range inputs {
+		t := makeRequest(input.transformID, input.sideInputID, input.tok)
+		tokens = append(tokens, t)
+	}
+
+	s.SetValidTokens(tokens...)
+	if len(s.idsToTokens) != len(inputs) {
+		t.Errorf("Missing tokens, expected %v, got %v", len(inputs), len(s.idsToTokens))
+	}
+
+	for i, input := range inputs {
+		// Check that the token is in the valid list
+		if !s.isValid(input.tok) {
+			t.Errorf("error in input %v, token %v is not valid", i, input.tok)
+		}
+		// Check that the mapping of IDs to tokens is correct
+		mapped := s.idsToTokens[input.transformID+input.sideInputID]
+		if mapped != input.tok {
+			t.Errorf("token mismatch for input %v, expected %v, got %v", i, input.tok, mapped)
+		}
+	}
+}
+
+func TestSetValidTokens_ClearingBetween(t *testing.T) {
+	inputs := []struct {
+		transformID string
+		sideInputID string
+		tk          token
+	}{
+		{
+			"t1",
+			"s1",
+			"tok1",
+		},
+		{
+			"t2",
+			"s2",
+			"tok2",
+		},
+		{
+			"t3",
+			"s3",
+			"tok3",
+		},
+	}
+
+	var s SideInputCache
+	err := s.Init(1)
+	if err != nil {
+		t.Fatalf("cache init failed, got %v", err)
+	}
+
+	for i, input := range inputs {
+		tok := makeRequest(input.transformID, input.sideInputID, input.tk)
+
+		s.SetValidTokens(tok)
+
+		// Check that the token is in the valid list
+		if !s.isValid(input.tk) {
+			t.Errorf("error in input %v, token %v is not valid", i, input.tk)
+		}
+		// Check that the mapping of IDs to tokens is correct
+		mapped := s.idsToTokens[input.transformID+input.sideInputID]
+		if mapped != input.tk {
+			t.Errorf("token mismatch for input %v, expected %v, got %v", i, input.tk, mapped)
+		}
+
+		s.CompleteBundle(tok)
+	}
+
+	for k, _ := range s.validTokens {
+		if s.validTokens[k] != 0 {
+			t.Errorf("token count mismatch for token %v, expected 0, got %v", k, s.validTokens[k])
+		}
+	}
+}
+
+func TestSetCache_Eviction(t *testing.T) {
+	var s SideInputCache
+	err := s.Init(1)
+	if err != nil {
+		t.Fatalf("cache init failed, got %v", err)
+	}
+
+	tokOne := makeRequest("t1", "s1", "tok1")
+	inOne := makeTestReusableInput("t1", "s1", 10)
+	s.SetValidTokens(tokOne)
+	s.SetCache("t1", "s1", inOne)
+	// Mark bundle as complete, drop count for tokOne to 0
+	s.CompleteBundle(tokOne)
+
+	tokTwo := makeRequest("t2", "s2", "tok2")
+	inTwo := makeTestReusableInput("t2", "s2", 20)
+	s.SetValidTokens(tokTwo)
+	s.SetCache("t2", "s2", inTwo)
+
+	if len(s.cache) != 1 {
+		t.Errorf("cache size incorrect, expected 1, got %v", len(s.cache))
+	}
+	if s.metrics.Evictions != 1 {
+		t.Errorf("number evictions incorrect, expected 1, got %v", s.metrics.Evictions)
+	}
+}
+
+func TestSetCache_EvictionFailure(t *testing.T) {
+	var s SideInputCache
+	err := s.Init(1)
+	if err != nil {
+		t.Fatalf("cache init failed, got %v", err)
+	}
+
+	tokOne := makeRequest("t1", "s1", "tok1")
+	inOne := makeTestReusableInput("t1", "s1", 10)
+
+	tokTwo := makeRequest("t2", "s2", "tok2")
+	inTwo := makeTestReusableInput("t2", "s2", 20)
+
+	s.SetValidTokens(tokOne, tokTwo)
+	s.SetCache("t1", "s1", inOne)
+	// Should fail to evict because the first token is still valid
+	s.SetCache("t2", "s2", inTwo)
+	// Cache should not exceed size 1
+	if len(s.cache) != 1 {
+		t.Errorf("cache size incorrect, expected 1, got %v", len(s.cache))
+	}
+	if s.metrics.InUseEvictions != 1 {
+		t.Errorf("number of failed evicition calls incorrect, expected 1, got %v", s.metrics.InUseEvictions)
+	}
+}
diff --git a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java
index 511f839..f4ab8bb 100644
--- a/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java
+++ b/sdks/java/fn-execution/src/main/java/org/apache/beam/sdk/fn/stream/DataStreams.java
@@ -156,16 +156,20 @@
   }
 
   /**
-   * An adapter which converts an {@link InputStream} to an {@link Iterator} of {@code T} values
-   * using the specified {@link Coder}.
+   * An adapter which converts an {@link InputStream} to a {@link PrefetchableIterator} of {@code T}
+   * values using the specified {@link Coder}.
    *
    * <p>Note that this adapter follows the Beam Fn API specification for forcing values that decode
    * consuming zero bytes to consuming exactly one byte.
    *
    * <p>Note that access to the underlying {@link InputStream} is lazy and will only be invoked on
-   * first access to {@link #next()} or {@link #hasNext()}.
+   * first access to {@link #next}, {@link #hasNext}, {@link #isReady}, and {@link #prefetch}.
+   *
+   * <p>Note that {@link #isReady} and {@link #prefetch} rely on non-empty {@link ByteString}s being
+   * returned via the underlying {@link PrefetchableIterator} otherwise the {@link #prefetch} will
+   * seemingly make zero progress yet will actually advance through the empty pages.
    */
-  public static class DataStreamDecoder<T> implements Iterator<T> {
+  public static class DataStreamDecoder<T> implements PrefetchableIterator<T> {
 
     private enum State {
       READ_REQUIRED,
@@ -173,13 +177,13 @@
       EOF
     }
 
-    private final Iterator<ByteString> inputByteStrings;
+    private final PrefetchableIterator<ByteString> inputByteStrings;
     private final Inbound inbound;
     private final Coder<T> coder;
     private State currentState;
     private T next;
 
-    public DataStreamDecoder(Coder<T> coder, Iterator<ByteString> inputStream) {
+    public DataStreamDecoder(Coder<T> coder, PrefetchableIterator<ByteString> inputStream) {
       this.currentState = State.READ_REQUIRED;
       this.coder = coder;
       this.inputByteStrings = inputStream;
@@ -187,6 +191,31 @@
     }
 
     @Override
+    public boolean isReady() {
+      switch (currentState) {
+        case EOF:
+          return true;
+        case READ_REQUIRED:
+          try {
+            return inbound.isReady();
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        case HAS_NEXT:
+          return true;
+        default:
+          throw new IllegalStateException(String.format("Unknown state %s", currentState));
+      }
+    }
+
+    @Override
+    public void prefetch() {
+      if (!isReady()) {
+        inputByteStrings.prefetch();
+      }
+    }
+
+    @Override
     public boolean hasNext() {
       switch (currentState) {
         case EOF:
@@ -232,8 +261,8 @@
     private static final InputStream EMPTY_STREAM = ByteString.EMPTY.newInput();
 
     /**
-     * An input stream which concatenates multiple {@link ByteString}s. Lazily accesses the first
-     * {@link Iterator} on first access of this input stream.
+     * An input stream which concatenates multiple {@link ByteString}s. Lazily accesses the {@link
+     * Iterator} on first access of this input stream.
      *
      * <p>Closing this input stream has no effect.
      */
@@ -245,6 +274,22 @@
         this.currentStream = EMPTY_STREAM;
       }
 
+      public boolean isReady() throws IOException {
+        // Note that ByteString#newInput is guaranteed to return the length of the entire ByteString
+        // minus the number of bytes that have been read so far and can be reliably used to tell
+        // us whether we are at the end of the stream.
+        while (currentStream.available() == 0) {
+          if (!inputByteStrings.isReady()) {
+            return false;
+          }
+          if (!inputByteStrings.hasNext()) {
+            return true;
+          }
+          currentStream = inputByteStrings.next().newInput();
+        }
+        return true;
+      }
+
       public boolean isEof() throws IOException {
         // Note that ByteString#newInput is guaranteed to return the length of the entire ByteString
         // minus the number of bytes that have been read so far and can be reliably used to tell
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java
index 9dd5ee4..a8b48e8 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/DataStreamsTest.java
@@ -23,6 +23,7 @@
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 
 import java.io.IOException;
@@ -106,7 +107,7 @@
     }
 
     @Test
-    public void testNonEmptyInputStreamWithZeroLengthCoder() throws Exception {
+    public void testNonEmptyInputStreamWithZeroLengthEncoding() throws Exception {
       CountingOutputStream countingOutputStream =
           new CountingOutputStream(ByteStreams.nullOutputStream());
       GlobalWindow.Coder.INSTANCE.encode(GlobalWindow.INSTANCE, countingOutputStream);
@@ -115,6 +116,55 @@
       testDecoderWith(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE, GlobalWindow.INSTANCE);
     }
 
+    @Test
+    public void testPrefetch() throws Exception {
+      List<ByteString> encodings = new ArrayList<>();
+      {
+        ByteString.Output encoding = ByteString.newOutput();
+        StringUtf8Coder.of().encode("A", encoding);
+        StringUtf8Coder.of().encode("BC", encoding);
+        encodings.add(encoding.toByteString());
+      }
+      encodings.add(ByteString.EMPTY);
+      {
+        ByteString.Output encoding = ByteString.newOutput();
+        StringUtf8Coder.of().encode("DEF", encoding);
+        StringUtf8Coder.of().encode("GHIJ", encoding);
+        encodings.add(encoding.toByteString());
+      }
+
+      PrefetchableIteratorsTest.ReadyAfterPrefetchUntilNext<ByteString> iterator =
+          new PrefetchableIteratorsTest.ReadyAfterPrefetchUntilNext<>(encodings.iterator());
+      PrefetchableIterator<String> decoder =
+          new DataStreamDecoder<>(StringUtf8Coder.of(), iterator);
+      assertFalse(decoder.isReady());
+      decoder.prefetch();
+      assertTrue(decoder.isReady());
+      assertEquals(1, iterator.getNumPrefetchCalls());
+
+      decoder.next();
+      // Now we will have moved off of the empty byte array that we start with so prefetch will
+      // do nothing since we are ready
+      assertTrue(decoder.isReady());
+      decoder.prefetch();
+      assertEquals(1, iterator.getNumPrefetchCalls());
+
+      decoder.next();
+      // Now we are at the end of the first ByteString so we expect a prefetch to pass through
+      assertFalse(decoder.isReady());
+      decoder.prefetch();
+      assertEquals(2, iterator.getNumPrefetchCalls());
+      // We also expect the decoder to not be ready since the next byte string is empty which
+      // would require us to move to the next page. This typically wouldn't happen in practice
+      // though because we expect non empty pages.
+      assertFalse(decoder.isReady());
+
+      // Prefetching will allow us to move to the third ByteString
+      decoder.prefetch();
+      assertEquals(3, iterator.getNumPrefetchCalls());
+      assertTrue(decoder.isReady());
+    }
+
     private <T> void testDecoderWith(Coder<T> coder, T... expected) throws IOException {
       ByteString.Output output = ByteString.newOutput();
       for (T value : expected) {
@@ -131,7 +181,9 @@
     }
 
     private <T> void testDecoderWith(Coder<T> coder, T[] expected, List<ByteString> encoded) {
-      Iterator<T> decoder = new DataStreamDecoder<>(coder, encoded.iterator());
+      Iterator<T> decoder =
+          new DataStreamDecoder<>(
+              coder, PrefetchableIterators.maybePrefetchable(encoded.iterator()));
 
       Object[] actual = Iterators.toArray(decoder, Object.class);
       assertArrayEquals(expected, actual);
diff --git a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/PrefetchableIteratorsTest.java b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/PrefetchableIteratorsTest.java
index 6131634..9ada175 100644
--- a/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/PrefetchableIteratorsTest.java
+++ b/sdks/java/fn-execution/src/test/java/org/apache/beam/sdk/fn/stream/PrefetchableIteratorsTest.java
@@ -120,10 +120,14 @@
         "F");
   }
 
-  private static class NeverReady implements PrefetchableIterator<String> {
-    PrefetchableIterator<String> delegate = PrefetchableIterators.fromArray("A", "B");
+  public static class NeverReady<T> implements PrefetchableIterator<T> {
+    private final Iterator<T> delegate;
     int prefetchCalled;
 
+    public NeverReady(Iterator<T> delegate) {
+      this.delegate = delegate;
+    }
+
     @Override
     public boolean isReady() {
       return false;
@@ -140,74 +144,117 @@
     }
 
     @Override
-    public String next() {
+    public T next() {
       return delegate.next();
     }
+
+    public int getNumPrefetchCalls() {
+      return prefetchCalled;
+    }
   }
 
-  private static class ReadyAfterPrefetch extends NeverReady {
+  public static class ReadyAfterPrefetch<T> extends NeverReady<T> {
+
+    public ReadyAfterPrefetch(Iterator<T> delegate) {
+      super(delegate);
+    }
+
     @Override
     public boolean isReady() {
       return prefetchCalled > 0;
     }
   }
 
+  public static class ReadyAfterPrefetchUntilNext<T> extends ReadyAfterPrefetch<T> {
+    boolean advancedSincePrefetch;
+
+    public ReadyAfterPrefetchUntilNext(Iterator<T> delegate) {
+      super(delegate);
+    }
+
+    @Override
+    public boolean isReady() {
+      return !advancedSincePrefetch && super.isReady();
+    }
+
+    @Override
+    public void prefetch() {
+      advancedSincePrefetch = false;
+      super.prefetch();
+    }
+
+    @Override
+    public T next() {
+      advancedSincePrefetch = true;
+      return super.next();
+    }
+
+    @Override
+    public boolean hasNext() {
+      advancedSincePrefetch = true;
+      return super.hasNext();
+    }
+  }
+
   @Test
   public void testConcatIsReadyAdvancesToNextIteratorWhenAble() {
-    NeverReady readyAfterPrefetch1 = new NeverReady();
-    ReadyAfterPrefetch readyAfterPrefetch2 = new ReadyAfterPrefetch();
-    ReadyAfterPrefetch readyAfterPrefetch3 = new ReadyAfterPrefetch();
+    NeverReady<String> readyAfterPrefetch1 =
+        new NeverReady<>(PrefetchableIterators.fromArray("A", "B"));
+    ReadyAfterPrefetch<String> readyAfterPrefetch2 =
+        new ReadyAfterPrefetch<>(PrefetchableIterators.fromArray("A", "B"));
+    ReadyAfterPrefetch<String> readyAfterPrefetch3 =
+        new ReadyAfterPrefetch<>(PrefetchableIterators.fromArray("A", "B"));
 
     PrefetchableIterator<String> iterator =
         PrefetchableIterators.concat(readyAfterPrefetch1, readyAfterPrefetch2, readyAfterPrefetch3);
 
     // Expect no prefetches yet
-    assertEquals(0, readyAfterPrefetch1.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch2.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch3.prefetchCalled);
+    assertEquals(0, readyAfterPrefetch1.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch2.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch3.getNumPrefetchCalls());
 
     // We expect to attempt to prefetch for the first time.
     iterator.prefetch();
-    assertEquals(1, readyAfterPrefetch1.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch2.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch3.prefetchCalled);
+    assertEquals(1, readyAfterPrefetch1.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch2.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch3.getNumPrefetchCalls());
     iterator.next();
 
     // We expect to attempt to prefetch again since we aren't ready.
     iterator.prefetch();
-    assertEquals(2, readyAfterPrefetch1.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch2.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch3.prefetchCalled);
+    assertEquals(2, readyAfterPrefetch1.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch2.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch3.getNumPrefetchCalls());
     iterator.next();
 
     // The current iterator is done but is never ready so we can't advance to the next one and
     // expect another prefetch to go to the current iterator.
     iterator.prefetch();
-    assertEquals(3, readyAfterPrefetch1.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch2.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch3.prefetchCalled);
+    assertEquals(3, readyAfterPrefetch1.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch2.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch3.getNumPrefetchCalls());
     iterator.next();
 
     // Now that we know the last iterator is done and have advanced to the next one we expect
     // prefetch to go through
     iterator.prefetch();
-    assertEquals(3, readyAfterPrefetch1.prefetchCalled);
-    assertEquals(1, readyAfterPrefetch2.prefetchCalled);
-    assertEquals(0, readyAfterPrefetch3.prefetchCalled);
+    assertEquals(3, readyAfterPrefetch1.getNumPrefetchCalls());
+    assertEquals(1, readyAfterPrefetch2.getNumPrefetchCalls());
+    assertEquals(0, readyAfterPrefetch3.getNumPrefetchCalls());
     iterator.next();
 
     // The last iterator is done so we should be able to prefetch the next one before advancing
     iterator.prefetch();
-    assertEquals(3, readyAfterPrefetch1.prefetchCalled);
-    assertEquals(1, readyAfterPrefetch2.prefetchCalled);
-    assertEquals(1, readyAfterPrefetch3.prefetchCalled);
+    assertEquals(3, readyAfterPrefetch1.getNumPrefetchCalls());
+    assertEquals(1, readyAfterPrefetch2.getNumPrefetchCalls());
+    assertEquals(1, readyAfterPrefetch3.getNumPrefetchCalls());
     iterator.next();
 
     // The current iterator is ready so no additional prefetch is necessary
     iterator.prefetch();
-    assertEquals(3, readyAfterPrefetch1.prefetchCalled);
-    assertEquals(1, readyAfterPrefetch2.prefetchCalled);
-    assertEquals(1, readyAfterPrefetch3.prefetchCalled);
+    assertEquals(3, readyAfterPrefetch1.getNumPrefetchCalls());
+    assertEquals(1, readyAfterPrefetch2.getNumPrefetchCalls());
+    assertEquals(1, readyAfterPrefetch3.getNumPrefetchCalls());
     iterator.next();
   }
 
diff --git a/sdks/java/harness/build.gradle b/sdks/java/harness/build.gradle
index 3c859ae..6337cd4 100644
--- a/sdks/java/harness/build.gradle
+++ b/sdks/java/harness/build.gradle
@@ -72,6 +72,7 @@
   testCompile library.java.mockito_core
   testCompile project(path: ":sdks:java:core", configuration: "shadowTest")
   testCompile project(":runners:core-construction-java")
+  testCompile project(path: ":sdks:java:fn-execution", configuration: "testRuntime")
   shadowTestRuntimeClasspath library.java.slf4j_jdk14
   jmhCompile project(path: ":sdks:java:harness", configuration: "shadowTest")
   jmhRuntime library.java.slf4j_jdk14
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java
index 5ddf0ae..777036a 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/BagUserState.java
@@ -26,6 +26,8 @@
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateClearRequest;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
 import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterable;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterables;
 import org.apache.beam.vendor.grpc.v1p36p0.com.google.protobuf.ByteString;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 
@@ -49,7 +51,7 @@
   private final BeamFnStateClient beamFnStateClient;
   private final StateRequest request;
   private final Coder<T> valueCoder;
-  private Iterable<T> oldValues;
+  private PrefetchableIterable<T> oldValues;
   private ArrayList<T> newValues;
   private boolean isClosed;
 
@@ -80,19 +82,19 @@
     this.newValues = new ArrayList<>();
   }
 
-  public Iterable<T> get() {
+  public PrefetchableIterable<T> get() {
     checkState(
         !isClosed,
         "Bag user state is no longer usable because it is closed for %s",
         request.getStateKey());
     if (oldValues == null) {
       // If we were cleared we should disregard old values.
-      return Iterables.limit(Collections.unmodifiableList(newValues), newValues.size());
+      return PrefetchableIterables.limit(Collections.unmodifiableList(newValues), newValues.size());
     } else if (newValues.isEmpty()) {
       // If we have no new values then just return the old values.
       return oldValues;
     }
-    return Iterables.concat(
+    return PrefetchableIterables.concat(
         oldValues, Iterables.limit(Collections.unmodifiableList(newValues), newValues.size()));
   }
 
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
index 2f2789e..5a931c5 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
@@ -38,7 +38,6 @@
 import org.apache.beam.sdk.state.MapState;
 import org.apache.beam.sdk.state.OrderedListState;
 import org.apache.beam.sdk.state.ReadableState;
-import org.apache.beam.sdk.state.ReadableStates;
 import org.apache.beam.sdk.state.SetState;
 import org.apache.beam.sdk.state.StateBinder;
 import org.apache.beam.sdk.state.StateContext;
@@ -264,7 +263,7 @@
 
                   @Override
                   public ValueState<T> readLater() {
-                    // TODO(BEAM-12802): Support prefetching.
+                    impl.get().iterator().prefetch();
                     return this;
                   }
                 };
@@ -310,7 +309,7 @@
 
                   @Override
                   public BagState<T> readLater() {
-                    // TODO(BEAM-12802): Support prefetching.
+                    impl.get().iterator().prefetch();
                     return this;
                   }
 
@@ -391,6 +390,7 @@
 
                   @Override
                   public CombiningState<ElementT, AccumT, ResultT> readLater() {
+                    impl.get().iterator().prefetch();
                     return this;
                   }
 
@@ -412,7 +412,18 @@
 
                   @Override
                   public ReadableState<Boolean> isEmpty() {
-                    return ReadableStates.immediate(!impl.get().iterator().hasNext());
+                    return new ReadableState<Boolean>() {
+                      @Override
+                      public @Nullable Boolean read() {
+                        return !impl.get().iterator().hasNext();
+                      }
+
+                      @Override
+                      public ReadableState<Boolean> readLater() {
+                        impl.get().iterator().prefetch();
+                        return this;
+                      }
+                    };
                   }
 
                   @Override
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterable.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterable.java
index cfc76cf..7828f93 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterable.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterable.java
@@ -21,6 +21,9 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.NoSuchElementException;
+import java.util.Objects;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterable;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterator;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
@@ -28,30 +31,42 @@
  * Converts an iterator to an iterable lazily loading values from the underlying iterator and
  * caching them to support reiteration.
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
-})
-class LazyCachingIteratorToIterable<T> implements Iterable<T> {
+class LazyCachingIteratorToIterable<T> implements PrefetchableIterable<T> {
   private final List<T> cachedElements;
-  private final Iterator<T> iterator;
+  private final PrefetchableIterator<T> iterator;
 
-  public LazyCachingIteratorToIterable(Iterator<T> iterator) {
+  public LazyCachingIteratorToIterable(PrefetchableIterator<T> iterator) {
     this.cachedElements = new ArrayList<>();
     this.iterator = iterator;
   }
 
   @Override
-  public Iterator<T> iterator() {
+  public PrefetchableIterator<T> iterator() {
     return new CachingIterator();
   }
 
   /** An {@link Iterator} which adds and fetched values into the cached elements list. */
-  private class CachingIterator implements Iterator<T> {
+  private class CachingIterator implements PrefetchableIterator<T> {
     private int position = 0;
 
     private CachingIterator() {}
 
     @Override
+    public boolean isReady() {
+      if (position < cachedElements.size()) {
+        return true;
+      }
+      return iterator.isReady();
+    }
+
+    @Override
+    public void prefetch() {
+      if (!isReady()) {
+        iterator.prefetch();
+      }
+    }
+
+    @Override
     public boolean hasNext() {
       // The order of the short circuit is important below.
       return position < cachedElements.size() || iterator.hasNext();
@@ -76,7 +91,7 @@
 
   @Override
   public int hashCode() {
-    return iterator.hasNext() ? iterator.next().hashCode() : -1789023489;
+    return iterator.hasNext() ? Objects.hashCode(iterator.next()) : -1789023489;
   }
 
   @Override
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
index 22be306..1026ba5 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
@@ -17,20 +17,21 @@
  */
 package org.apache.beam.fn.harness.state;
 
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
-import java.util.function.Supplier;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateGetRequest;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse;
 import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.fn.stream.DataStreams;
+import org.apache.beam.sdk.fn.stream.DataStreams.DataStreamDecoder;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterable;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterator;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterators;
 import org.apache.beam.vendor.grpc.v1p36p0.com.google.protobuf.ByteString;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Throwables;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 
 /**
  * Adapters which convert a a logical series of chunks using continuation tokens over the Beam Fn
@@ -54,7 +55,7 @@
    *     only) chunk of a state stream. This state request will be populated with a continuation
    *     token to request further chunks of the stream if required.
    */
-  public static Iterator<ByteString> readAllStartingFrom(
+  public static PrefetchableIterator<ByteString> readAllStartingFrom(
       BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk) {
     return new LazyBlockingStateFetchingIterator(beamFnStateClient, stateRequestForFirstChunk);
   }
@@ -74,94 +75,142 @@
    *     token to request further chunks of the stream if required.
    * @param valueCoder A coder for decoding the state stream.
    */
-  public static <T> Iterable<T> readAllAndDecodeStartingFrom(
+  public static <T> PrefetchableIterable<T> readAllAndDecodeStartingFrom(
       BeamFnStateClient beamFnStateClient,
       StateRequest stateRequestForFirstChunk,
       Coder<T> valueCoder) {
-    FirstPageAndRemainder firstPageAndRemainder =
-        new FirstPageAndRemainder(beamFnStateClient, stateRequestForFirstChunk);
-    return Iterables.concat(
-        new LazyCachingIteratorToIterable<T>(
-            new DataStreams.DataStreamDecoder<>(
-                valueCoder, new LazySingletonIterator<>(firstPageAndRemainder::firstPage))),
-        () -> new DataStreams.DataStreamDecoder<>(valueCoder, firstPageAndRemainder.remainder()));
-  }
-
-  /** A iterable that contains a single element, provided by a Supplier which is invoked lazily. */
-  static class LazySingletonIterator<T> implements Iterator<T> {
-
-    private final Supplier<T> supplier;
-    private boolean hasNext;
-
-    private LazySingletonIterator(Supplier<T> supplier) {
-      this.supplier = supplier;
-      hasNext = true;
-    }
-
-    @Override
-    public boolean hasNext() {
-      return hasNext;
-    }
-
-    @Override
-    public T next() {
-      hasNext = false;
-      return supplier.get();
-    }
+    return new FirstPageAndRemainder<>(beamFnStateClient, stateRequestForFirstChunk, valueCoder);
   }
 
   /**
-   * An helper class that (lazily) gives the first page of a paginated state request separately from
+   * A helper class that (lazily) gives the first page of a paginated state request separately from
    * all the remaining pages.
    */
-  static class FirstPageAndRemainder {
+  @VisibleForTesting
+  static class FirstPageAndRemainder<T> implements PrefetchableIterable<T> {
     private final BeamFnStateClient beamFnStateClient;
     private final StateRequest stateRequestForFirstChunk;
-    private ByteString firstPage = null;
+    private final Coder<T> valueCoder;
+    private LazyCachingIteratorToIterable<T> firstPage;
+    private CompletableFuture<StateResponse> firstPageResponseFuture;
     private ByteString continuationToken;
 
-    private FirstPageAndRemainder(
-        BeamFnStateClient beamFnStateClient, StateRequest stateRequestForFirstChunk) {
+    FirstPageAndRemainder(
+        BeamFnStateClient beamFnStateClient,
+        StateRequest stateRequestForFirstChunk,
+        Coder<T> valueCoder) {
       this.beamFnStateClient = beamFnStateClient;
       this.stateRequestForFirstChunk = stateRequestForFirstChunk;
+      this.valueCoder = valueCoder;
     }
 
-    public ByteString firstPage() {
-      if (firstPage == null) {
-        CompletableFuture<StateResponse> stateResponseFuture = new CompletableFuture<>();
+    @Override
+    public PrefetchableIterator<T> iterator() {
+      return new PrefetchableIterator<T>() {
+        PrefetchableIterator<T> delegate;
+
+        private void ensureDelegateExists() {
+          if (delegate == null) {
+            // Fetch the first page if necessary
+            prefetchFirstPage();
+            if (firstPage == null) {
+              StateResponse stateResponse;
+              try {
+                stateResponse = firstPageResponseFuture.get();
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new IllegalStateException(e);
+              } catch (ExecutionException e) {
+                if (e.getCause() == null) {
+                  throw new IllegalStateException(e);
+                }
+                Throwables.throwIfUnchecked(e.getCause());
+                throw new IllegalStateException(e.getCause());
+              }
+              continuationToken = stateResponse.getGet().getContinuationToken();
+              firstPage =
+                  new LazyCachingIteratorToIterable<>(
+                      new DataStreamDecoder<>(
+                          valueCoder,
+                          PrefetchableIterators.fromArray(stateResponse.getGet().getData())));
+            }
+
+            if (ByteString.EMPTY.equals((continuationToken))) {
+              delegate = firstPage.iterator();
+            } else {
+              delegate =
+                  PrefetchableIterators.concat(
+                      firstPage.iterator(),
+                      new DataStreamDecoder<>(
+                          valueCoder,
+                          new LazyBlockingStateFetchingIterator(
+                              beamFnStateClient,
+                              stateRequestForFirstChunk
+                                  .toBuilder()
+                                  .setGet(
+                                      StateGetRequest.newBuilder()
+                                          .setContinuationToken(continuationToken))
+                                  .build())));
+            }
+          }
+        }
+
+        @Override
+        public boolean isReady() {
+          if (delegate == null) {
+            if (firstPageResponseFuture != null) {
+              return firstPageResponseFuture.isDone();
+            }
+            return false;
+          }
+          return delegate.isReady();
+        }
+
+        @Override
+        public void prefetch() {
+          if (firstPageResponseFuture == null) {
+            prefetchFirstPage();
+          } else if (delegate != null && !delegate.isReady()) {
+            delegate.prefetch();
+          }
+        }
+
+        @Override
+        public boolean hasNext() {
+          if (delegate == null) {
+            // Ensure that we prefetch the second page after the first has been accessed.
+            // Prefetching subsequent pages after the first will be handled by the
+            // LazyBlockingStateFetchingIterator
+            ensureDelegateExists();
+            boolean rval = delegate.hasNext();
+            delegate.prefetch();
+            return rval;
+          }
+          return delegate.hasNext();
+        }
+
+        @Override
+        public T next() {
+          if (delegate == null) {
+            // Ensure that we prefetch the second page after the first has been accessed.
+            // Prefetching subsequent pages after the first will be handled by the
+            // LazyBlockingStateFetchingIterator
+            ensureDelegateExists();
+            T rval = delegate.next();
+            delegate.prefetch();
+            return rval;
+          }
+          return delegate.next();
+        }
+      };
+    }
+
+    private void prefetchFirstPage() {
+      if (firstPageResponseFuture == null) {
+        firstPageResponseFuture = new CompletableFuture<>();
         beamFnStateClient.handle(
             stateRequestForFirstChunk.toBuilder().setGet(stateRequestForFirstChunk.getGet()),
-            stateResponseFuture);
-        StateResponse stateResponse;
-        try {
-          stateResponse = stateResponseFuture.get();
-        } catch (InterruptedException e) {
-          Thread.currentThread().interrupt();
-          throw new IllegalStateException(e);
-        } catch (ExecutionException e) {
-          if (e.getCause() == null) {
-            throw new IllegalStateException(e);
-          }
-          Throwables.throwIfUnchecked(e.getCause());
-          throw new IllegalStateException(e.getCause());
-        }
-        continuationToken = stateResponse.getGet().getContinuationToken();
-        firstPage = stateResponse.getGet().getData();
-      }
-      return firstPage;
-    }
-
-    public Iterator<ByteString> remainder() {
-      firstPage();
-      if (ByteString.EMPTY.equals(continuationToken)) {
-        return Collections.emptyIterator();
-      } else {
-        return new LazyBlockingStateFetchingIterator(
-            beamFnStateClient,
-            stateRequestForFirstChunk
-                .toBuilder()
-                .setGet(StateGetRequest.newBuilder().setContinuationToken(continuationToken))
-                .build());
+            firstPageResponseFuture);
       }
     }
   }
@@ -169,10 +218,11 @@
   /**
    * An {@link Iterator} which fetches {@link ByteString} chunks using the State API.
    *
-   * <p>This iterator will only request a chunk on first access. Subsiquently it eagerly pre-fetches
-   * one future chunks at a time.
+   * <p>This iterator will only request a chunk on first access. Subsequently it eagerly pre-fetches
+   * one future chunk at a time.
    */
-  static class LazyBlockingStateFetchingIterator implements Iterator<ByteString> {
+  @VisibleForTesting
+  static class LazyBlockingStateFetchingIterator implements PrefetchableIterator<ByteString> {
 
     private enum State {
       READ_REQUIRED,
@@ -195,8 +245,17 @@
       this.continuationToken = stateRequestForFirstChunk.getGet().getContinuationToken();
     }
 
-    private void prefetch() {
-      if (prefetchedResponse == null && currentState == State.READ_REQUIRED) {
+    @Override
+    public boolean isReady() {
+      if (prefetchedResponse == null) {
+        return currentState != State.READ_REQUIRED;
+      }
+      return prefetchedResponse.isDone();
+    }
+
+    @Override
+    public void prefetch() {
+      if (currentState == State.READ_REQUIRED && prefetchedResponse == null) {
         prefetchedResponse = new CompletableFuture<>();
         beamFnStateClient.handle(
             stateRequestForFirstChunk
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterableTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterableTest.java
index 7597128..0914b01 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterableTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/LazyCachingIteratorToIterableTest.java
@@ -25,8 +25,11 @@
 
 import java.util.Iterator;
 import java.util.NoSuchElementException;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterable;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterator;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterators;
+import org.apache.beam.sdk.fn.stream.PrefetchableIteratorsTest.ReadyAfterPrefetch;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -40,7 +43,8 @@
 
   @Test
   public void testEmptyIterator() {
-    Iterable<Object> iterable = new LazyCachingIteratorToIterable<>(Iterators.forArray());
+    Iterable<Object> iterable =
+        new LazyCachingIteratorToIterable<>(PrefetchableIterators.emptyIterator());
     assertArrayEquals(new Object[0], Iterables.toArray(iterable, Object.class));
     // iterate multiple times
     assertArrayEquals(new Object[0], Iterables.toArray(iterable, Object.class));
@@ -52,7 +56,7 @@
   @Test
   public void testInterleavedIteration() {
     Iterable<String> iterable =
-        new LazyCachingIteratorToIterable<>(Iterators.forArray("A", "B", "C"));
+        new LazyCachingIteratorToIterable<>(PrefetchableIterators.fromArray("A", "B", "C"));
 
     Iterator<String> iterator1 = iterable.iterator();
     assertTrue(iterator1.hasNext());
@@ -77,14 +81,45 @@
 
   @Test
   public void testEqualsAndHashCode() {
-    Iterable<String> iterA = new LazyCachingIteratorToIterable<>(Iterators.forArray("A", "B", "C"));
-    Iterable<String> iterB = new LazyCachingIteratorToIterable<>(Iterators.forArray("A", "B", "C"));
-    Iterable<String> iterC = new LazyCachingIteratorToIterable<>(Iterators.forArray());
-    Iterable<String> iterD = new LazyCachingIteratorToIterable<>(Iterators.forArray());
+    Iterable<String> iterA =
+        new LazyCachingIteratorToIterable<>(PrefetchableIterators.fromArray("A", "B", "C"));
+    Iterable<String> iterB =
+        new LazyCachingIteratorToIterable<>(PrefetchableIterators.fromArray("A", "B", "C"));
+    Iterable<String> iterC = new LazyCachingIteratorToIterable<>(PrefetchableIterators.fromArray());
+    Iterable<String> iterD = new LazyCachingIteratorToIterable<>(PrefetchableIterators.fromArray());
     assertEquals(iterA, iterB);
     assertEquals(iterC, iterD);
     assertNotEquals(iterA, iterC);
     assertEquals(iterA.hashCode(), iterB.hashCode());
     assertEquals(iterC.hashCode(), iterD.hashCode());
   }
+
+  @Test
+  public void testPrefetch() {
+    ReadyAfterPrefetch<String> underlying =
+        new ReadyAfterPrefetch<>(PrefetchableIterators.fromArray("A", "B", "C"));
+    PrefetchableIterable<String> iterable = new LazyCachingIteratorToIterable<>(underlying);
+    PrefetchableIterator<String> iterator1 = iterable.iterator();
+    PrefetchableIterator<String> iterator2 = iterable.iterator();
+
+    // Check that the lazy iterable doesn't do any prefetch/access on instantiation
+    assertFalse(underlying.isReady());
+    assertFalse(iterator1.isReady());
+    assertFalse(iterator2.isReady());
+
+    // Check that if both iterators prefetch there is only one prefetch for the underlying iterator
+    // iterator.
+    iterator1.prefetch();
+    assertEquals(1, underlying.getNumPrefetchCalls());
+    iterator2.prefetch();
+    assertEquals(1, underlying.getNumPrefetchCalls());
+
+    // Check that if that one iterator has advanced, the second doesn't perform any prefetch since
+    // the element is now cached.
+    iterator1.next();
+    iterator1.next();
+    iterator2.next();
+    iterator2.prefetch();
+    assertEquals(1, underlying.getNumPrefetchCalls());
+  }
 }
diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java
index fc729cc..384d2df 100644
--- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java
+++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/StateFetchingIteratorsTest.java
@@ -19,12 +19,16 @@
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
+import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Iterator;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
+import org.apache.beam.fn.harness.state.StateFetchingIterators.FirstPageAndRemainder;
 import org.apache.beam.fn.harness.state.StateFetchingIterators.LazyBlockingStateFetchingIterator;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateGetResponse;
 import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest;
@@ -32,16 +36,56 @@
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.CoderException;
 import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterable;
+import org.apache.beam.sdk.fn.stream.PrefetchableIterator;
 import org.apache.beam.sdk.util.CoderUtils;
 import org.apache.beam.vendor.grpc.v1p36p0.com.google.protobuf.ByteString;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterators;
 import org.junit.Test;
+import org.junit.experimental.runners.Enclosed;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 /** Tests for {@link StateFetchingIterators}. */
+@RunWith(Enclosed.class)
 public class StateFetchingIteratorsTest {
+
+  private static BeamFnStateClient fakeStateClient(
+      AtomicInteger callCount, ByteString... expected) {
+    return (requestBuilder, response) -> {
+      callCount.incrementAndGet();
+      if (expected.length == 0) {
+        response.complete(
+            StateResponse.newBuilder()
+                .setId(requestBuilder.getId())
+                .setGet(StateGetResponse.newBuilder())
+                .build());
+        return;
+      }
+
+      ByteString continuationToken = requestBuilder.getGet().getContinuationToken();
+
+      int requestedPosition = 0; // Default position is 0
+      if (!ByteString.EMPTY.equals(continuationToken)) {
+        requestedPosition = Integer.parseInt(continuationToken.toStringUtf8());
+      }
+
+      // Compute the new continuation token
+      ByteString newContinuationToken = ByteString.EMPTY;
+      if (requestedPosition != expected.length - 1) {
+        newContinuationToken = ByteString.copyFromUtf8(Integer.toString(requestedPosition + 1));
+      }
+      response.complete(
+          StateResponse.newBuilder()
+              .setId(requestBuilder.getId())
+              .setGet(
+                  StateGetResponse.newBuilder()
+                      .setData(expected[requestedPosition])
+                      .setContinuationToken(newContinuationToken))
+              .build());
+    };
+  }
+
   /** Tests for {@link StateFetchingIterators.LazyBlockingStateFetchingIterator}. */
   @RunWith(JUnit4.class)
   public static class LazyBlockingStateFetchingIteratorTest {
@@ -77,49 +121,55 @@
           ByteString.EMPTY);
     }
 
-    private BeamFnStateClient fakeStateClient(AtomicInteger callCount, ByteString... expected) {
-      return (requestBuilder, response) -> {
-        callCount.incrementAndGet();
-        if (expected.length == 0) {
-          response.complete(
-              StateResponse.newBuilder()
-                  .setId(requestBuilder.getId())
-                  .setGet(StateGetResponse.newBuilder())
-                  .build());
-          return;
-        }
-
-        ByteString continuationToken = requestBuilder.getGet().getContinuationToken();
-
-        int requestedPosition = 0; // Default position is 0
-        if (!ByteString.EMPTY.equals(continuationToken)) {
-          requestedPosition = Integer.parseInt(continuationToken.toStringUtf8());
-        }
-
-        // Compute the new continuation token
-        ByteString newContinuationToken = ByteString.EMPTY;
-        if (requestedPosition != expected.length - 1) {
-          newContinuationToken = ByteString.copyFromUtf8(Integer.toString(requestedPosition + 1));
-        }
-        response.complete(
-            StateResponse.newBuilder()
-                .setId(requestBuilder.getId())
-                .setGet(
-                    StateGetResponse.newBuilder()
-                        .setData(expected[requestedPosition])
-                        .setContinuationToken(newContinuationToken))
-                .build());
-      };
+    @Test
+    public void testPrefetchIgnoredWhenExistingPrefetchOngoing() throws Exception {
+      AtomicInteger callCount = new AtomicInteger();
+      BeamFnStateClient fakeStateClient =
+          new BeamFnStateClient() {
+            @Override
+            public void handle(
+                StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> response) {
+              callCount.incrementAndGet();
+            }
+          };
+      PrefetchableIterator<ByteString> byteStrings =
+          new LazyBlockingStateFetchingIterator(fakeStateClient, StateRequest.getDefaultInstance());
+      assertEquals(0, callCount.get());
+      byteStrings.prefetch();
+      assertEquals(1, callCount.get()); // first prefetch
+      byteStrings.prefetch();
+      assertEquals(1, callCount.get()); // subsequent is ignored
     }
 
     private void testFetch(ByteString... expected) {
       AtomicInteger callCount = new AtomicInteger();
       BeamFnStateClient fakeStateClient = fakeStateClient(callCount, expected);
-      Iterator<ByteString> byteStrings =
+      PrefetchableIterator<ByteString> byteStrings =
           new LazyBlockingStateFetchingIterator(fakeStateClient, StateRequest.getDefaultInstance());
       assertEquals(0, callCount.get()); // Ensure it's fully lazy.
-      assertArrayEquals(expected, Iterators.toArray(byteStrings, Object.class));
+      assertFalse(byteStrings.isReady());
+
+      // Prefetch every second element in the iterator capturing the results
+      List<ByteString> results = new ArrayList<>();
+      for (int i = 0; i < expected.length; ++i) {
+        if (i % 2 == 0) {
+          // Ensure that prefetch performs the call
+          byteStrings.prefetch();
+          assertEquals(i + 1, callCount.get());
+          assertTrue(byteStrings.isReady());
+        }
+        assertTrue(byteStrings.hasNext());
+        results.add(byteStrings.next());
+      }
+      assertFalse(byteStrings.hasNext());
+      assertTrue(byteStrings.isReady());
+
+      assertEquals(Arrays.asList(expected), results);
     }
+  }
+
+  @RunWith(JUnit4.class)
+  public static class FirstPageAndRemainderTest {
 
     @Test
     public void testEmptyValues() throws Exception {
@@ -133,7 +183,7 @@
 
     @Test
     public void testManyValues() throws Exception {
-      testFetchValues(VarIntCoder.of(), 11, 37, 389, 5077);
+      testFetchValues(VarIntCoder.of(), 1, 22, 333, 4444, 55555, 666666);
     }
 
     private <T> void testFetchValues(Coder<T> coder, T... expected) {
@@ -153,35 +203,42 @@
       AtomicInteger callCount = new AtomicInteger();
       BeamFnStateClient fakeStateClient =
           fakeStateClient(callCount, Iterables.toArray(byteStrings, ByteString.class));
-      Iterable<T> values =
-          StateFetchingIterators.readAllAndDecodeStartingFrom(
-              fakeStateClient, StateRequest.getDefaultInstance(), coder);
+      PrefetchableIterable<T> values =
+          new FirstPageAndRemainder<>(fakeStateClient, StateRequest.getDefaultInstance(), coder);
 
       // Ensure it's fully lazy.
       assertEquals(0, callCount.get());
-      Iterator<T> valuesIter = values.iterator();
+      PrefetchableIterator<T> valuesIter = values.iterator();
+      assertFalse(valuesIter.isReady());
       assertEquals(0, callCount.get());
 
-      // No more is read than necessary.
-      if (valuesIter.hasNext()) {
-        valuesIter.next();
-      }
+      // Ensure that the first page result is cached across multiple iterators and subsequent
+      // iterators are ready and prefetch does nothing
+      valuesIter.prefetch();
+      assertTrue(valuesIter.isReady());
       assertEquals(1, callCount.get());
 
-      // The first page is cached.
-      Iterator<T> valuesIter2 = values.iterator();
-      assertEquals(1, callCount.get());
-      if (valuesIter2.hasNext()) {
-        valuesIter2.next();
-      }
+      PrefetchableIterator<T> valuesIter2 = values.iterator();
+      assertTrue(valuesIter2.isReady());
+      valuesIter2.prefetch();
       assertEquals(1, callCount.get());
 
-      if (valuesIter.hasNext()) {
-        valuesIter.next();
-        // Subsequent pages are pre-fetched, so after accessing the second page,
-        // the third should be requested.
-        assertEquals(3, callCount.get());
+      // Prefetch every second element in the iterator capturing the results
+      List<T> results = new ArrayList<>();
+      for (int i = 0; i < expected.length; ++i) {
+        if (i % 2 == 1) {
+          // Ensure that prefetch performs the call
+          valuesIter2.prefetch();
+          assertTrue(valuesIter2.isReady());
+          // Note that this is i+2 because we expect to prefetch the page after the current one
+          // We also have to bound it to the max number of pages
+          assertEquals(Math.min(i + 2, expected.length), callCount.get());
+        }
+        assertTrue(valuesIter2.hasNext());
+        results.add(valuesIter2.next());
       }
+      assertFalse(valuesIter2.hasNext());
+      assertTrue(valuesIter2.isReady());
 
       // The contents agree.
       assertArrayEquals(expected, Iterables.toArray(values, Object.class));
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java
index dd5202e..74d6636 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1.java
@@ -1218,6 +1218,10 @@
         List<Cursor> cursors = new ArrayList<>(partitionQueryResponse.getPartitionsList());
         cursors.sort(CURSOR_REFERENCE_VALUE_COMPARATOR);
         final int size = cursors.size();
+        if (size == 0) {
+          emit(c, dbRoot, structuredQuery.toBuilder());
+          return;
+        }
         final int lastIdx = size - 1;
         for (int i = 0; i < size; i++) {
           Cursor curr = cursors.get(i);
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessages.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessages.java
index b0cc681..bf6a288 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessages.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsub/PubsubMessages.java
@@ -24,25 +24,36 @@
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
 
 /** Common util functions for converting between PubsubMessage proto and {@link PubsubMessage}. */
-public class PubsubMessages {
+public final class PubsubMessages {
+  private PubsubMessages() {}
+
+  public static com.google.pubsub.v1.PubsubMessage toProto(PubsubMessage input) {
+    Map<String, String> attributes = input.getAttributeMap();
+    com.google.pubsub.v1.PubsubMessage.Builder message =
+        com.google.pubsub.v1.PubsubMessage.newBuilder()
+            .setData(ByteString.copyFrom(input.getPayload()));
+    // TODO(BEAM-8085) this should not be null
+    if (attributes != null) {
+      message.putAllAttributes(attributes);
+    }
+    String messageId = input.getMessageId();
+    if (messageId != null) {
+      message.setMessageId(messageId);
+    }
+    return message.build();
+  }
+
+  public static PubsubMessage fromProto(com.google.pubsub.v1.PubsubMessage input) {
+    return new PubsubMessage(
+        input.getData().toByteArray(), input.getAttributesMap(), input.getMessageId());
+  }
+
   // Convert the PubsubMessage to a PubsubMessage proto, then return its serialized representation.
   public static class ParsePayloadAsPubsubMessageProto
       implements SerializableFunction<PubsubMessage, byte[]> {
     @Override
     public byte[] apply(PubsubMessage input) {
-      Map<String, String> attributes = input.getAttributeMap();
-      com.google.pubsub.v1.PubsubMessage.Builder message =
-          com.google.pubsub.v1.PubsubMessage.newBuilder()
-              .setData(ByteString.copyFrom(input.getPayload()));
-      // TODO(BEAM-8085) this should not be null
-      if (attributes != null) {
-        message.putAllAttributes(attributes);
-      }
-      String messageId = input.getMessageId();
-      if (messageId != null) {
-        message.setMessageId(messageId);
-      }
-      return message.build().toByteArray();
+      return toProto(input).toByteArray();
     }
   }
 
@@ -54,8 +65,7 @@
       try {
         com.google.pubsub.v1.PubsubMessage message =
             com.google.pubsub.v1.PubsubMessage.parseFrom(input);
-        return new PubsubMessage(
-            message.getData().toByteArray(), message.getAttributesMap(), message.getMessageId());
+        return fromProto(message);
       } catch (InvalidProtocolBufferException e) {
         throw new RuntimeException("Could not decode Pubsub message", e);
       }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/CloudPubsubChecks.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/CloudPubsubChecks.java
deleted file mode 100644
index 6dc1516..0000000
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/CloudPubsubChecks.java
+++ /dev/null
@@ -1,51 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package org.apache.beam.sdk.io.gcp.pubsublite;
-
-import static com.google.cloud.pubsublite.cloudpubsub.MessageTransforms.toCpsPublishTransformer;
-
-import com.google.cloud.pubsublite.Message;
-import com.google.cloud.pubsublite.proto.PubSubMessage;
-import org.apache.beam.sdk.transforms.MapElements;
-import org.apache.beam.sdk.transforms.PTransform;
-import org.apache.beam.sdk.values.PCollection;
-import org.apache.beam.sdk.values.TypeDescriptor;
-
-/**
- * A class providing a conversion validity check between Cloud Pub/Sub and Pub/Sub Lite message
- * types.
- */
-public final class CloudPubsubChecks {
-  private CloudPubsubChecks() {}
-
-  /**
-   * Ensure that all messages that pass through can be converted to Cloud Pub/Sub messages using the
-   * standard transformation methods in the client library.
-   *
-   * <p>Will fail the pipeline if a message has multiple attributes per key.
-   */
-  public static PTransform<PCollection<? extends PubSubMessage>, PCollection<PubSubMessage>>
-      ensureUsableAsCloudPubsub() {
-    return MapElements.into(TypeDescriptor.of(PubSubMessage.class))
-        .via(
-            message -> {
-              Object unused = toCpsPublishTransformer().transform(Message.fromProto(message));
-              return message;
-            });
-  }
-}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/CloudPubsubTransforms.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/CloudPubsubTransforms.java
new file mode 100644
index 0000000..1140c11
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/CloudPubsubTransforms.java
@@ -0,0 +1,104 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsublite;
+
+import static com.google.cloud.pubsublite.cloudpubsub.MessageTransforms.fromCpsPublishTransformer;
+import static com.google.cloud.pubsublite.cloudpubsub.MessageTransforms.toCpsPublishTransformer;
+import static com.google.cloud.pubsublite.cloudpubsub.MessageTransforms.toCpsSubscribeTransformer;
+
+import com.google.cloud.pubsublite.Message;
+import com.google.cloud.pubsublite.cloudpubsub.KeyExtractor;
+import com.google.cloud.pubsublite.proto.PubSubMessage;
+import com.google.cloud.pubsublite.proto.SequencedMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessage;
+import org.apache.beam.sdk.io.gcp.pubsub.PubsubMessages;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TypeDescriptor;
+
+/** A class providing transforms between Cloud Pub/Sub and Pub/Sub Lite message types. */
+public final class CloudPubsubTransforms {
+  private CloudPubsubTransforms() {}
+  /**
+   * Ensure that all messages that pass through can be converted to Cloud Pub/Sub messages using the
+   * standard transformation methods in the client library.
+   *
+   * <p>Will fail the pipeline if a message has multiple attributes per key.
+   */
+  public static PTransform<PCollection<PubSubMessage>, PCollection<PubSubMessage>>
+      ensureUsableAsCloudPubsub() {
+    return new PTransform<PCollection<PubSubMessage>, PCollection<PubSubMessage>>() {
+      @Override
+      public PCollection<PubSubMessage> expand(PCollection<PubSubMessage> input) {
+        return input.apply(
+            MapElements.into(TypeDescriptor.of(PubSubMessage.class))
+                .via(
+                    message -> {
+                      Object unused =
+                          toCpsPublishTransformer().transform(Message.fromProto(message));
+                      return message;
+                    }));
+      }
+    };
+  }
+
+  /**
+   * Transform messages read from Pub/Sub Lite to their equivalent Cloud Pub/Sub Message that would
+   * have been read from PubsubIO.
+   *
+   * <p>Will fail the pipeline if a message has multiple attributes per map key.
+   */
+  public static PTransform<PCollection<SequencedMessage>, PCollection<PubsubMessage>>
+      toCloudPubsubMessages() {
+    return new PTransform<PCollection<SequencedMessage>, PCollection<PubsubMessage>>() {
+      @Override
+      public PCollection<PubsubMessage> expand(PCollection<SequencedMessage> input) {
+        return input.apply(
+            MapElements.into(TypeDescriptor.of(PubsubMessage.class))
+                .via(
+                    message ->
+                        PubsubMessages.fromProto(
+                            toCpsSubscribeTransformer()
+                                .transform(
+                                    com.google.cloud.pubsublite.SequencedMessage.fromProto(
+                                        message)))));
+      }
+    };
+  }
+
+  /**
+   * Transform messages publishable using PubsubIO to their equivalent Pub/Sub Lite publishable
+   * message.
+   */
+  public static PTransform<PCollection<PubsubMessage>, PCollection<PubSubMessage>>
+      fromCloudPubsubMessages() {
+    return new PTransform<PCollection<PubsubMessage>, PCollection<PubSubMessage>>() {
+      @Override
+      public PCollection<PubSubMessage> expand(PCollection<PubsubMessage> input) {
+        return input.apply(
+            MapElements.into(TypeDescriptor.of(PubSubMessage.class))
+                .via(
+                    message ->
+                        fromCpsPublishTransformer(KeyExtractor.DEFAULT)
+                            .transform(PubsubMessages.toProto(message))
+                            .toProto()));
+      }
+    };
+  }
+}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/ManagedBacklogReaderFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/ManagedBacklogReaderFactory.java
new file mode 100644
index 0000000..de0cf43
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/ManagedBacklogReaderFactory.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsublite;
+
+import java.io.Serializable;
+
+/**
+ * A ManagedBacklogReaderFactory produces TopicBacklogReaders and tears down any produced readers
+ * when it is itself closed.
+ *
+ * <p>close() should never be called on produced readers.
+ */
+public interface ManagedBacklogReaderFactory extends AutoCloseable, Serializable {
+  TopicBacklogReader newReader(SubscriptionPartition subscriptionPartition);
+
+  @Override
+  void close();
+}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/ManagedBacklogReaderFactoryImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/ManagedBacklogReaderFactoryImpl.java
new file mode 100644
index 0000000..9a337bf
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/ManagedBacklogReaderFactoryImpl.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsublite;
+
+import com.google.api.gax.rpc.ApiException;
+import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.proto.ComputeMessageStatsResponse;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.concurrent.GuardedBy;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+
+public class ManagedBacklogReaderFactoryImpl implements ManagedBacklogReaderFactory {
+  private final SerializableFunction<SubscriptionPartition, TopicBacklogReader> newReader;
+
+  @GuardedBy("this")
+  private final Map<SubscriptionPartition, TopicBacklogReader> readers = new HashMap<>();
+
+  ManagedBacklogReaderFactoryImpl(
+      SerializableFunction<SubscriptionPartition, TopicBacklogReader> newReader) {
+    this.newReader = newReader;
+  }
+
+  private static final class NonCloseableTopicBacklogReader implements TopicBacklogReader {
+    private final TopicBacklogReader underlying;
+
+    NonCloseableTopicBacklogReader(TopicBacklogReader underlying) {
+      this.underlying = underlying;
+    }
+
+    @Override
+    public ComputeMessageStatsResponse computeMessageStats(Offset offset) throws ApiException {
+      return underlying.computeMessageStats(offset);
+    }
+
+    @Override
+    public void close() {
+      throw new IllegalArgumentException(
+          "Cannot call close() on a reader returned from ManagedBacklogReaderFactory.");
+    }
+  }
+
+  @Override
+  public synchronized TopicBacklogReader newReader(SubscriptionPartition subscriptionPartition) {
+    return new NonCloseableTopicBacklogReader(
+        readers.computeIfAbsent(subscriptionPartition, newReader::apply));
+  }
+
+  @Override
+  public synchronized void close() {
+    readers.values().forEach(TopicBacklogReader::close);
+  }
+}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRange.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRange.java
new file mode 100644
index 0000000..b39d87e
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRange.java
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsublite;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.sdk.coders.DefaultCoder;
+import org.apache.beam.sdk.io.range.OffsetRange;
+
+@AutoValue
+@DefaultCoder(OffsetByteRangeCoder.class)
+abstract class OffsetByteRange {
+  abstract OffsetRange getRange();
+
+  abstract long getByteCount();
+
+  static OffsetByteRange of(OffsetRange range, long byteCount) {
+    return new AutoValue_OffsetByteRange(range, byteCount);
+  }
+
+  static OffsetByteRange of(OffsetRange range) {
+    return of(range, 0);
+  }
+}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeCoder.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeCoder.java
new file mode 100644
index 0000000..076cda1
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeCoder.java
@@ -0,0 +1,63 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsublite;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderProvider;
+import org.apache.beam.sdk.coders.CoderProviders;
+import org.apache.beam.sdk.coders.DelegateCoder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.beam.sdk.io.range.OffsetRange;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TypeDescriptor;
+
+public class OffsetByteRangeCoder extends AtomicCoder<OffsetByteRange> {
+  private static final Coder<OffsetByteRange> CODER =
+      DelegateCoder.of(
+          KvCoder.of(OffsetRange.Coder.of(), VarLongCoder.of()),
+          OffsetByteRangeCoder::toKv,
+          OffsetByteRangeCoder::fromKv);
+
+  private static KV<OffsetRange, Long> toKv(OffsetByteRange value) {
+    return KV.of(value.getRange(), value.getByteCount());
+  }
+
+  private static OffsetByteRange fromKv(KV<OffsetRange, Long> kv) {
+    return OffsetByteRange.of(kv.getKey(), kv.getValue());
+  }
+
+  @Override
+  public void encode(OffsetByteRange value, OutputStream outStream) throws IOException {
+    CODER.encode(value, outStream);
+  }
+
+  @Override
+  public OffsetByteRange decode(InputStream inStream) throws IOException {
+    return CODER.decode(inStream);
+  }
+
+  public static CoderProvider getCoderProvider() {
+    return CoderProviders.forCoder(
+        TypeDescriptor.of(OffsetByteRange.class), new OffsetByteRangeCoder());
+  }
+}
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTracker.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTracker.java
index 608af8f..da9aaaa 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTracker.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTracker.java
@@ -26,8 +26,6 @@
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.io.range.OffsetRange;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
-import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
 import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Stopwatch;
 import org.joda.time.Duration;
@@ -44,24 +42,27 @@
  * received. IMPORTANT: minTrackingTime must be strictly smaller than the SDF read timeout when it
  * would return ProcessContinuation.resume().
  */
-class OffsetByteRangeTracker extends RestrictionTracker<OffsetRange, OffsetByteProgress>
-    implements HasProgress {
-  private final TopicBacklogReader backlogReader;
+class OffsetByteRangeTracker extends TrackerWithProgress {
+  private final TopicBacklogReader unownedBacklogReader;
   private final Duration minTrackingTime;
   private final long minBytesReceived;
   private final Stopwatch stopwatch;
-  private OffsetRange range;
+  private OffsetByteRange range;
   private @Nullable Long lastClaimed;
-  private long byteCount = 0;
 
   public OffsetByteRangeTracker(
-      OffsetRange range,
-      TopicBacklogReader backlogReader,
+      OffsetByteRange range,
+      TopicBacklogReader unownedBacklogReader,
       Stopwatch stopwatch,
       Duration minTrackingTime,
       long minBytesReceived) {
-    checkArgument(range.getTo() == Long.MAX_VALUE);
-    this.backlogReader = backlogReader;
+    checkArgument(
+        range.getRange().getTo() == Long.MAX_VALUE,
+        "May only construct OffsetByteRangeTracker with an unbounded range with no progress.");
+    checkArgument(
+        range.getByteCount() == 0L,
+        "May only construct OffsetByteRangeTracker with an unbounded range with no progress.");
+    this.unownedBacklogReader = unownedBacklogReader;
     this.minTrackingTime = minTrackingTime;
     this.minBytesReceived = minBytesReceived;
     this.stopwatch = stopwatch.reset().start();
@@ -69,11 +70,6 @@
   }
 
   @Override
-  public void finalize() {
-    this.backlogReader.close();
-  }
-
-  @Override
   public IsBounded isBounded() {
     return IsBounded.UNBOUNDED;
   }
@@ -87,32 +83,32 @@
         position.lastOffset().value(),
         lastClaimed);
     checkArgument(
-        toClaim >= range.getFrom(),
+        toClaim >= range.getRange().getFrom(),
         "Trying to claim offset %s before start of the range %s",
         toClaim,
         range);
     // split() has already been called, truncating this range. No more offsets may be claimed.
-    if (range.getTo() != Long.MAX_VALUE) {
-      boolean isRangeEmpty = range.getTo() == range.getFrom();
-      boolean isValidClosedRange = nextOffset() == range.getTo();
+    if (range.getRange().getTo() != Long.MAX_VALUE) {
+      boolean isRangeEmpty = range.getRange().getTo() == range.getRange().getFrom();
+      boolean isValidClosedRange = nextOffset() == range.getRange().getTo();
       checkState(
           isRangeEmpty || isValidClosedRange,
           "Violated class precondition: offset range improperly split. Please report a beam bug.");
       return false;
     }
     lastClaimed = toClaim;
-    byteCount += position.batchBytes();
+    range = OffsetByteRange.of(range.getRange(), range.getByteCount() + position.batchBytes());
     return true;
   }
 
   @Override
-  public OffsetRange currentRestriction() {
+  public OffsetByteRange currentRestriction() {
     return range;
   }
 
   private long nextOffset() {
     checkState(lastClaimed == null || lastClaimed < Long.MAX_VALUE);
-    return lastClaimed == null ? currentRestriction().getFrom() : lastClaimed + 1;
+    return lastClaimed == null ? currentRestriction().getRange().getFrom() : lastClaimed + 1;
   }
 
   /**
@@ -124,29 +120,33 @@
     if (duration.isLongerThan(minTrackingTime)) {
       return true;
     }
-    if (byteCount >= minBytesReceived) {
+    if (currentRestriction().getByteCount() >= minBytesReceived) {
       return true;
     }
     return false;
   }
 
   @Override
-  public @Nullable SplitResult<OffsetRange> trySplit(double fractionOfRemainder) {
+  public @Nullable SplitResult<OffsetByteRange> trySplit(double fractionOfRemainder) {
     // Cannot split a bounded range. This should already be completely claimed.
-    if (range.getTo() != Long.MAX_VALUE) {
+    if (range.getRange().getTo() != Long.MAX_VALUE) {
       return null;
     }
     if (!receivedEnough()) {
       return null;
     }
-    range = new OffsetRange(currentRestriction().getFrom(), nextOffset());
-    return SplitResult.of(this.range, new OffsetRange(nextOffset(), Long.MAX_VALUE));
+    range =
+        OffsetByteRange.of(
+            new OffsetRange(currentRestriction().getRange().getFrom(), nextOffset()),
+            range.getByteCount());
+    return SplitResult.of(
+        this.range, OffsetByteRange.of(new OffsetRange(nextOffset(), Long.MAX_VALUE), 0));
   }
 
   @Override
   @SuppressWarnings("unboxing.of.nullable")
   public void checkDone() throws IllegalStateException {
-    if (range.getFrom() == range.getTo()) {
+    if (range.getRange().getFrom() == range.getRange().getTo()) {
       return;
     }
     checkState(
@@ -155,18 +155,18 @@
         range);
     long lastClaimedNotNull = checkNotNull(lastClaimed);
     checkState(
-        lastClaimedNotNull >= range.getTo() - 1,
+        lastClaimedNotNull >= range.getRange().getTo() - 1,
         "Last attempted offset was %s in range %s, claiming work in [%s, %s) was not attempted",
         lastClaimedNotNull,
         range,
         lastClaimedNotNull + 1,
-        range.getTo());
+        range.getRange().getTo());
   }
 
   @Override
   public Progress getProgress() {
     ComputeMessageStatsResponse stats =
-        this.backlogReader.computeMessageStats(Offset.of(nextOffset()));
-    return Progress.from(byteCount, stats.getMessageBytes());
+        this.unownedBacklogReader.computeMessageStats(Offset.of(nextOffset()));
+    return Progress.from(range.getByteCount(), stats.getMessageBytes());
   }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerServerPublisherCache.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerServerPublisherCache.java
index 623e20c..d7526d8 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerServerPublisherCache.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerServerPublisherCache.java
@@ -27,4 +27,8 @@
   private PerServerPublisherCache() {}
 
   static final PublisherCache PUBLISHER_CACHE = new PublisherCache();
+
+  static {
+    Runtime.getRuntime().addShutdownHook(new Thread(PUBLISHER_CACHE::close));
+  }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdf.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdf.java
index a9f7a43..fdf7920 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdf.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdf.java
@@ -17,13 +17,12 @@
  */
 package org.apache.beam.sdk.io.gcp.pubsublite;
 
-import static com.google.cloud.pubsublite.internal.ExtractStatus.toCanonical;
-import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+import static com.google.cloud.pubsublite.internal.wire.ApiServiceUtils.blockingShutdown;
 
 import com.google.cloud.pubsublite.Offset;
+import com.google.cloud.pubsublite.internal.ExtractStatus;
 import com.google.cloud.pubsublite.internal.wire.Committer;
 import com.google.cloud.pubsublite.proto.SequencedMessage;
-import java.util.concurrent.ExecutionException;
 import org.apache.beam.sdk.io.range.OffsetRange;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.SerializableBiFunction;
@@ -35,31 +34,35 @@
 
 class PerSubscriptionPartitionSdf extends DoFn<SubscriptionPartition, SequencedMessage> {
   private final Duration maxSleepTime;
+  private final ManagedBacklogReaderFactory backlogReaderFactory;
   private final SubscriptionPartitionProcessorFactory processorFactory;
   private final SerializableFunction<SubscriptionPartition, InitialOffsetReader>
       offsetReaderFactory;
-  private final SerializableBiFunction<
-          SubscriptionPartition, OffsetRange, RestrictionTracker<OffsetRange, OffsetByteProgress>>
+  private final SerializableBiFunction<TopicBacklogReader, OffsetByteRange, TrackerWithProgress>
       trackerFactory;
   private final SerializableFunction<SubscriptionPartition, Committer> committerFactory;
 
   PerSubscriptionPartitionSdf(
       Duration maxSleepTime,
+      ManagedBacklogReaderFactory backlogReaderFactory,
       SerializableFunction<SubscriptionPartition, InitialOffsetReader> offsetReaderFactory,
-      SerializableBiFunction<
-              SubscriptionPartition,
-              OffsetRange,
-              RestrictionTracker<OffsetRange, OffsetByteProgress>>
+      SerializableBiFunction<TopicBacklogReader, OffsetByteRange, TrackerWithProgress>
           trackerFactory,
       SubscriptionPartitionProcessorFactory processorFactory,
       SerializableFunction<SubscriptionPartition, Committer> committerFactory) {
     this.maxSleepTime = maxSleepTime;
+    this.backlogReaderFactory = backlogReaderFactory;
     this.processorFactory = processorFactory;
     this.offsetReaderFactory = offsetReaderFactory;
     this.trackerFactory = trackerFactory;
     this.committerFactory = committerFactory;
   }
 
+  @Teardown
+  public void teardown() {
+    backlogReaderFactory.close();
+  }
+
   @GetInitialWatermarkEstimatorState
   public Instant getInitialWatermarkState() {
     return Instant.EPOCH;
@@ -72,7 +75,7 @@
 
   @ProcessElement
   public ProcessContinuation processElement(
-      RestrictionTracker<OffsetRange, OffsetByteProgress> tracker,
+      RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker,
       @Element SubscriptionPartition subscriptionPartition,
       OutputReceiver<SequencedMessage> receiver)
       throws Exception {
@@ -83,38 +86,44 @@
       processor
           .lastClaimed()
           .ifPresent(
-              lastClaimedOffset ->
-              /* TODO(boyuanzz): When default dataflow can use finalizers, undo this.
-              finalizer.afterBundleCommit(
-                  Instant.ofEpochMilli(Long.MAX_VALUE),
-                  () -> */ {
+              lastClaimedOffset -> {
                 Committer committer = committerFactory.apply(subscriptionPartition);
                 committer.startAsync().awaitRunning();
                 // Commit the next-to-deliver offset.
                 try {
                   committer.commitOffset(Offset.of(lastClaimedOffset.value() + 1)).get();
-                } catch (ExecutionException e) {
-                  throw toCanonical(checkArgumentNotNull(e.getCause())).underlying;
                 } catch (Exception e) {
-                  throw toCanonical(e).underlying;
+                  throw ExtractStatus.toCanonical(e).underlying;
                 }
-                committer.stopAsync().awaitTerminated();
+                blockingShutdown(committer);
               });
       return result;
     }
   }
 
   @GetInitialRestriction
-  public OffsetRange getInitialRestriction(@Element SubscriptionPartition subscriptionPartition) {
+  public OffsetByteRange getInitialRestriction(
+      @Element SubscriptionPartition subscriptionPartition) {
     try (InitialOffsetReader reader = offsetReaderFactory.apply(subscriptionPartition)) {
       Offset offset = reader.read();
-      return new OffsetRange(offset.value(), Long.MAX_VALUE /* open interval */);
+      return OffsetByteRange.of(
+          new OffsetRange(offset.value(), Long.MAX_VALUE /* open interval */));
     }
   }
 
   @NewTracker
-  public RestrictionTracker<OffsetRange, OffsetByteProgress> newTracker(
-      @Element SubscriptionPartition subscriptionPartition, @Restriction OffsetRange range) {
-    return trackerFactory.apply(subscriptionPartition, range);
+  public TrackerWithProgress newTracker(
+      @Element SubscriptionPartition subscriptionPartition, @Restriction OffsetByteRange range) {
+    return trackerFactory.apply(backlogReaderFactory.newReader(subscriptionPartition), range);
+  }
+
+  @GetSize
+  public double getSize(
+      @Element SubscriptionPartition subscriptionPartition,
+      @Restriction OffsetByteRange restriction) {
+    if (restriction.getRange().getTo() != Long.MAX_VALUE) {
+      return restriction.getByteCount();
+    }
+    return newTracker(subscriptionPartition, restriction).getProgress().getWorkRemaining();
   }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PublisherCache.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PublisherCache.java
index f8dc24b..3dbdec6 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PublisherCache.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PublisherCache.java
@@ -23,52 +23,50 @@
 import com.google.api.core.ApiService.State;
 import com.google.api.gax.rpc.ApiException;
 import com.google.cloud.pubsublite.MessageMetadata;
-import com.google.cloud.pubsublite.internal.CloseableMonitor;
 import com.google.cloud.pubsublite.internal.Publisher;
+import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import java.util.HashMap;
-import java.util.concurrent.Executor;
-import java.util.concurrent.Executors;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
 
 /** A map of working publishers by PublisherOptions. */
-class PublisherCache {
-  private final CloseableMonitor monitor = new CloseableMonitor();
-
-  private final Executor listenerExecutor = Executors.newSingleThreadExecutor();
-
-  @GuardedBy("monitor.monitor")
+class PublisherCache implements AutoCloseable {
+  @GuardedBy("this")
   private final HashMap<PublisherOptions, Publisher<MessageMetadata>> livePublishers =
       new HashMap<>();
 
-  Publisher<MessageMetadata> get(PublisherOptions options) throws ApiException {
+  private synchronized void evict(PublisherOptions options) {
+    livePublishers.remove(options);
+  }
+
+  synchronized Publisher<MessageMetadata> get(PublisherOptions options) throws ApiException {
     checkArgument(options.usesCache());
-    try (CloseableMonitor.Hold h = monitor.enter()) {
-      Publisher<MessageMetadata> publisher = livePublishers.get(options);
-      if (publisher != null) {
-        return publisher;
-      }
-      publisher = Publishers.newPublisher(options);
-      livePublishers.put(options, publisher);
-      publisher.addListener(
-          new Listener() {
-            @Override
-            public void failed(State s, Throwable t) {
-              try (CloseableMonitor.Hold h = monitor.enter()) {
-                livePublishers.remove(options);
-              }
-            }
-          },
-          listenerExecutor);
-      publisher.startAsync().awaitRunning();
+    Publisher<MessageMetadata> publisher = livePublishers.get(options);
+    if (publisher != null) {
       return publisher;
     }
+    publisher = Publishers.newPublisher(options);
+    livePublishers.put(options, publisher);
+    publisher.addListener(
+        new Listener() {
+          @Override
+          public void failed(State s, Throwable t) {
+            evict(options);
+          }
+        },
+        SystemExecutors.getFuturesExecutor());
+    publisher.startAsync().awaitRunning();
+    return publisher;
   }
 
   @VisibleForTesting
-  void set(PublisherOptions options, Publisher<MessageMetadata> toCache) {
-    try (CloseableMonitor.Hold h = monitor.enter()) {
-      livePublishers.put(options, toCache);
-    }
+  synchronized void set(PublisherOptions options, Publisher<MessageMetadata> toCache) {
+    livePublishers.put(options, toCache);
+  }
+
+  @Override
+  public synchronized void close() {
+    livePublishers.forEach(((options, publisher) -> publisher.stopAsync()));
+    livePublishers.clear();
   }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/Publishers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/Publishers.java
index 34012f7..67ea6cf 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/Publishers.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/Publishers.java
@@ -17,17 +17,27 @@
  */
 package org.apache.beam.sdk.io.gcp.pubsublite;
 
+import static com.google.cloud.pubsublite.internal.ExtractStatus.toCanonical;
 import static com.google.cloud.pubsublite.internal.UncheckedApiPreconditions.checkArgument;
+import static com.google.cloud.pubsublite.internal.wire.ServiceClients.addDefaultMetadata;
+import static com.google.cloud.pubsublite.internal.wire.ServiceClients.addDefaultSettings;
 
 import com.google.api.gax.rpc.ApiException;
 import com.google.cloud.pubsublite.AdminClient;
 import com.google.cloud.pubsublite.AdminClientSettings;
 import com.google.cloud.pubsublite.MessageMetadata;
-import com.google.cloud.pubsublite.TopicPath;
+import com.google.cloud.pubsublite.Partition;
+import com.google.cloud.pubsublite.cloudpubsub.PublisherSettings;
 import com.google.cloud.pubsublite.internal.Publisher;
 import com.google.cloud.pubsublite.internal.wire.PartitionCountWatchingPublisherSettings;
+import com.google.cloud.pubsublite.internal.wire.PubsubContext;
 import com.google.cloud.pubsublite.internal.wire.PubsubContext.Framework;
+import com.google.cloud.pubsublite.internal.wire.RoutingMetadata;
 import com.google.cloud.pubsublite.internal.wire.SinglePartitionPublisherBuilder;
+import com.google.cloud.pubsublite.v1.AdminServiceClient;
+import com.google.cloud.pubsublite.v1.AdminServiceSettings;
+import com.google.cloud.pubsublite.v1.PublisherServiceClient;
+import com.google.cloud.pubsublite.v1.PublisherServiceSettings;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.reflect.TypeToken;
 
 class Publishers {
@@ -35,6 +45,38 @@
 
   private Publishers() {}
 
+  private static AdminClient newAdminClient(PublisherOptions options) throws ApiException {
+    try {
+      return AdminClient.create(
+          AdminClientSettings.newBuilder()
+              .setServiceClient(
+                  AdminServiceClient.create(
+                      addDefaultSettings(
+                          options.topicPath().location().extractRegion(),
+                          AdminServiceSettings.newBuilder())))
+              .setRegion(options.topicPath().location().extractRegion())
+              .build());
+    } catch (Throwable t) {
+      throw toCanonical(t).underlying;
+    }
+  }
+
+  private static PublisherServiceClient newServiceClient(
+      PublisherOptions options, Partition partition) {
+    PublisherServiceSettings.Builder settingsBuilder = PublisherServiceSettings.newBuilder();
+    settingsBuilder =
+        addDefaultMetadata(
+            PubsubContext.of(FRAMEWORK),
+            RoutingMetadata.of(options.topicPath(), partition),
+            settingsBuilder);
+    try {
+      return PublisherServiceClient.create(
+          addDefaultSettings(options.topicPath().location().extractRegion(), settingsBuilder));
+    } catch (Throwable t) {
+      throw toCanonical(t).underlying;
+    }
+  }
+
   @SuppressWarnings("unchecked")
   static Publisher<MessageMetadata> newPublisher(PublisherOptions options) throws ApiException {
     SerializableSupplier<Object> supplier = options.publisherSupplier();
@@ -44,20 +86,18 @@
       checkArgument(token.isSupertypeOf(supplied.getClass()));
       return (Publisher<MessageMetadata>) supplied;
     }
-
-    TopicPath topic = options.topicPath();
-    PartitionCountWatchingPublisherSettings.Builder publisherSettings =
-        PartitionCountWatchingPublisherSettings.newBuilder()
-            .setTopic(topic)
-            .setPublisherFactory(
-                partition ->
-                    SinglePartitionPublisherBuilder.newBuilder()
-                        .setTopic(topic)
-                        .setPartition(partition)
-                        .build())
-            .setAdminClient(
-                AdminClient.create(
-                    AdminClientSettings.newBuilder().setRegion(topic.location().region()).build()));
-    return publisherSettings.build().instantiate();
+    return PartitionCountWatchingPublisherSettings.newBuilder()
+        .setTopic(options.topicPath())
+        .setPublisherFactory(
+            partition ->
+                SinglePartitionPublisherBuilder.newBuilder()
+                    .setTopic(options.topicPath())
+                    .setPartition(partition)
+                    .setServiceClient(newServiceClient(options, partition))
+                    .setBatchingSettings(PublisherSettings.DEFAULT_BATCHING_SETTINGS)
+                    .build())
+        .setAdminClient(newAdminClient(options))
+        .build()
+        .instantiate();
   }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteIO.java
index ca1f2be..b93ac61 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteIO.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteIO.java
@@ -107,7 +107,7 @@
    * }</pre>
    */
   public static PTransform<PCollection<PubSubMessage>, PDone> write(PublisherOptions options) {
-    return new PTransform<PCollection<PubSubMessage>, PDone>("PubsubLiteIO") {
+    return new PTransform<PCollection<PubSubMessage>, PDone>() {
       @Override
       public PDone expand(PCollection<PubSubMessage> input) {
         PubsubLiteSink sink = new PubsubLiteSink(options);
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteSink.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteSink.java
index d3acdfa..d0e3afa 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteSink.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/PubsubLiteSink.java
@@ -28,16 +28,14 @@
 import com.google.cloud.pubsublite.internal.CheckedApiException;
 import com.google.cloud.pubsublite.internal.ExtractStatus;
 import com.google.cloud.pubsublite.internal.Publisher;
+import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
 import com.google.cloud.pubsublite.proto.PubSubMessage;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
 import java.util.ArrayDeque;
 import java.util.Deque;
-import java.util.concurrent.Executor;
-import java.util.concurrent.Executors;
 import java.util.function.Consumer;
 import org.apache.beam.sdk.io.gcp.pubsublite.PublisherOrError.Kind;
 import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.MoreExecutors;
 
 /** A sink which publishes messages to Pub/Sub Lite. */
 @SuppressWarnings({
@@ -56,8 +54,6 @@
   @GuardedBy("this")
   private transient Deque<CheckedApiException> errorsSinceLastFinish;
 
-  private static final Executor executor = Executors.newCachedThreadPool();
-
   PubsubLiteSink(PublisherOptions options) {
     this.options = options;
   }
@@ -89,7 +85,7 @@
             onFailure.accept(t);
           }
         },
-        MoreExecutors.directExecutor());
+        SystemExecutors.getFuturesExecutor());
     if (!options.usesCache()) {
       publisher.startAsync();
     }
@@ -130,7 +126,7 @@
             onFailure.accept(t);
           }
         },
-        executor);
+        SystemExecutors.getFuturesExecutor());
   }
 
   // Intentionally don't flush on bundle finish to allow multi-sink client reuse.
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscribeTransform.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscribeTransform.java
index 9875880..b6a9f5d 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscribeTransform.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscribeTransform.java
@@ -23,6 +23,7 @@
 import com.google.api.gax.rpc.ApiException;
 import com.google.cloud.pubsublite.AdminClient;
 import com.google.cloud.pubsublite.AdminClientSettings;
+import com.google.cloud.pubsublite.Offset;
 import com.google.cloud.pubsublite.Partition;
 import com.google.cloud.pubsublite.TopicPath;
 import com.google.cloud.pubsublite.internal.wire.Committer;
@@ -31,8 +32,6 @@
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
-import org.apache.beam.sdk.io.range.OffsetRange;
-import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -54,10 +53,11 @@
     checkArgument(subscriptionPartition.subscription().equals(options.subscriptionPath()));
   }
 
-  private Subscriber newSubscriber(Partition partition, Consumer<List<SequencedMessage>> consumer) {
+  private Subscriber newSubscriber(
+      Partition partition, Offset initialOffset, Consumer<List<SequencedMessage>> consumer) {
     try {
       return options
-          .getSubscriberFactory(partition)
+          .getSubscriberFactory(partition, initialOffset)
           .newSubscriber(
               messages ->
                   consumer.accept(
@@ -71,23 +71,31 @@
 
   private SubscriptionPartitionProcessor newPartitionProcessor(
       SubscriptionPartition subscriptionPartition,
-      RestrictionTracker<OffsetRange, OffsetByteProgress> tracker,
+      RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker,
       OutputReceiver<SequencedMessage> receiver)
       throws ApiException {
     checkSubscription(subscriptionPartition);
     return new SubscriptionPartitionProcessorImpl(
         tracker,
         receiver,
-        consumer -> newSubscriber(subscriptionPartition.partition(), consumer),
+        consumer ->
+            newSubscriber(
+                subscriptionPartition.partition(),
+                Offset.of(tracker.currentRestriction().getRange().getFrom()),
+                consumer),
         options.flowControlSettings());
   }
 
-  private RestrictionTracker<OffsetRange, OffsetByteProgress> newRestrictionTracker(
-      SubscriptionPartition subscriptionPartition, OffsetRange initial) {
+  private TopicBacklogReader newBacklogReader(SubscriptionPartition subscriptionPartition) {
     checkSubscription(subscriptionPartition);
+    return options.getBacklogReader(subscriptionPartition.partition());
+  }
+
+  private TrackerWithProgress newRestrictionTracker(
+      TopicBacklogReader backlogReader, OffsetByteRange initial) {
     return new OffsetByteRangeTracker(
         initial,
-        options.getBacklogReader(subscriptionPartition.partition()),
+        backlogReader,
         Stopwatch.createUnstarted(),
         options.minBundleTimeout(),
         LongMath.saturatedMultiply(options.flowControlSettings().bytesOutstanding(), 10));
@@ -107,7 +115,7 @@
     try (AdminClient admin =
         AdminClient.create(
             AdminClientSettings.newBuilder()
-                .setRegion(options.subscriptionPath().location().region())
+                .setRegion(options.subscriptionPath().location().extractRegion())
                 .build())) {
       return TopicPath.parse(admin.getSubscription(options.subscriptionPath()).get().getTopic());
     } catch (Throwable t) {
@@ -118,25 +126,15 @@
   @Override
   public PCollection<SequencedMessage> expand(PBegin input) {
     PCollection<SubscriptionPartition> subscriptionPartitions;
-    if (options.partitions().isEmpty()) {
-      subscriptionPartitions =
-          input.apply(new SubscriptionPartitionLoader(getTopicPath(), options.subscriptionPath()));
-    } else {
-      subscriptionPartitions =
-          input.apply(
-              Create.of(
-                  options.partitions().stream()
-                      .map(
-                          partition ->
-                              SubscriptionPartition.of(options.subscriptionPath(), partition))
-                      .collect(Collectors.toList())));
-    }
+    subscriptionPartitions =
+        input.apply(new SubscriptionPartitionLoader(getTopicPath(), options.subscriptionPath()));
 
     return subscriptionPartitions.apply(
         ParDo.of(
             new PerSubscriptionPartitionSdf(
                 // Ensure we read for at least 5 seconds more than the bundle timeout.
                 options.minBundleTimeout().plus(Duration.standardSeconds(5)),
+                new ManagedBacklogReaderFactoryImpl(this::newBacklogReader),
                 this::newInitialOffsetReader,
                 this::newRestrictionTracker,
                 this::newPartitionProcessor,
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriberOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriberOptions.java
index 0d3afe2..a9625be 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriberOptions.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriberOptions.java
@@ -23,6 +23,7 @@
 
 import com.google.api.gax.rpc.ApiException;
 import com.google.auto.value.AutoValue;
+import com.google.cloud.pubsublite.Offset;
 import com.google.cloud.pubsublite.Partition;
 import com.google.cloud.pubsublite.SubscriptionPath;
 import com.google.cloud.pubsublite.cloudpubsub.FlowControlSettings;
@@ -35,13 +36,13 @@
 import com.google.cloud.pubsublite.internal.wire.RoutingMetadata;
 import com.google.cloud.pubsublite.internal.wire.SubscriberBuilder;
 import com.google.cloud.pubsublite.internal.wire.SubscriberFactory;
+import com.google.cloud.pubsublite.proto.Cursor;
+import com.google.cloud.pubsublite.proto.SeekRequest;
 import com.google.cloud.pubsublite.v1.CursorServiceClient;
 import com.google.cloud.pubsublite.v1.CursorServiceSettings;
 import com.google.cloud.pubsublite.v1.SubscriberServiceClient;
 import com.google.cloud.pubsublite.v1.SubscriberServiceSettings;
 import java.io.Serializable;
-import java.util.Set;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
 
@@ -69,11 +70,6 @@
   public abstract FlowControlSettings flowControlSettings();
 
   /**
-   * A set of partitions. If empty, continuously poll the set of partitions using an admin client.
-   */
-  public abstract Set<Partition> partitions();
-
-  /**
    * The minimum wall time to pass before allowing bundle closure.
    *
    * <p>Setting this to too small of a value will result in increased compute costs and lower
@@ -108,7 +104,6 @@
   public static Builder newBuilder() {
     Builder builder = new AutoValue_SubscriberOptions.Builder();
     return builder
-        .setPartitions(ImmutableSet.of())
         .setFlowControlSettings(DEFAULT_FLOW_CONTROL)
         .setMinBundleTimeout(MIN_BUNDLE_TIMEOUT);
   }
@@ -119,20 +114,19 @@
       throws ApiException {
     try {
       SubscriberServiceSettings.Builder settingsBuilder = SubscriberServiceSettings.newBuilder();
-
       settingsBuilder =
           addDefaultMetadata(
               PubsubContext.of(FRAMEWORK),
               RoutingMetadata.of(subscriptionPath(), partition),
               settingsBuilder);
       return SubscriberServiceClient.create(
-          addDefaultSettings(subscriptionPath().location().region(), settingsBuilder));
+          addDefaultSettings(subscriptionPath().location().extractRegion(), settingsBuilder));
     } catch (Throwable t) {
       throw toCanonical(t).underlying;
     }
   }
 
-  SubscriberFactory getSubscriberFactory(Partition partition) {
+  SubscriberFactory getSubscriberFactory(Partition partition, Offset initialOffset) {
     SubscriberFactory factory = subscriberFactory();
     if (factory != null) {
       return factory;
@@ -143,6 +137,10 @@
             .setSubscriptionPath(subscriptionPath())
             .setPartition(partition)
             .setServiceClient(newSubscriberServiceClient(partition))
+            .setInitialLocation(
+                SeekRequest.newBuilder()
+                    .setCursor(Cursor.newBuilder().setOffset(initialOffset.value()))
+                    .build())
             .build();
   }
 
@@ -150,7 +148,7 @@
     try {
       return CursorServiceClient.create(
           addDefaultSettings(
-              subscriptionPath().location().region(), CursorServiceSettings.newBuilder()));
+              subscriptionPath().location().extractRegion(), CursorServiceSettings.newBuilder()));
     } catch (Throwable t) {
       throw toCanonical(t).underlying;
     }
@@ -189,7 +187,7 @@
     return new InitialOffsetReaderImpl(
         CursorClient.create(
             CursorClientSettings.newBuilder()
-                .setRegion(subscriptionPath().location().region())
+                .setRegion(subscriptionPath().location().extractRegion())
                 .build()),
         subscriptionPath(),
         partition);
@@ -201,8 +199,6 @@
     public abstract Builder setSubscriptionPath(SubscriptionPath path);
 
     // Optional parameters
-    public abstract Builder setPartitions(Set<Partition> partitions);
-
     public abstract Builder setFlowControlSettings(FlowControlSettings flowControlSettings);
 
     public abstract Builder setMinBundleTimeout(Duration minBundleTimeout);
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionLoader.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionLoader.java
index 866e922..e411d80 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionLoader.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionLoader.java
@@ -92,7 +92,9 @@
                     })
                 .withPollInterval(pollDuration)
                 .withTerminationPerInput(
-                    terminate ? Watch.Growth.afterIterations(10) : Watch.Growth.never()));
+                    terminate
+                        ? Watch.Growth.afterTotalOf(pollDuration.multipliedBy(10))
+                        : Watch.Growth.never()));
     return partitions.apply(
         MapElements.into(TypeDescriptor.of(SubscriptionPartition.class))
             .via(kv -> SubscriptionPartition.of(subscription, kv.getValue())));
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorFactory.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorFactory.java
index 6bf3623..530c180 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorFactory.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorFactory.java
@@ -19,7 +19,6 @@
 
 import com.google.cloud.pubsublite.proto.SequencedMessage;
 import java.io.Serializable;
-import org.apache.beam.sdk.io.range.OffsetRange;
 import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
 import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
 
@@ -28,6 +27,6 @@
 
   SubscriptionPartitionProcessor newProcessor(
       SubscriptionPartition subscriptionPartition,
-      RestrictionTracker<OffsetRange, OffsetByteProgress> tracker,
+      RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker,
       OutputReceiver<SequencedMessage> receiver);
 }
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImpl.java
index 8d2a137..a086d18 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImpl.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImpl.java
@@ -17,6 +17,8 @@
  */
 package org.apache.beam.sdk.io.gcp.pubsublite;
 
+import static com.google.cloud.pubsublite.internal.wire.ApiServiceUtils.blockingShutdown;
+
 import com.google.api.core.ApiService.Listener;
 import com.google.api.core.ApiService.State;
 import com.google.cloud.pubsublite.Offset;
@@ -24,9 +26,8 @@
 import com.google.cloud.pubsublite.internal.CheckedApiException;
 import com.google.cloud.pubsublite.internal.ExtractStatus;
 import com.google.cloud.pubsublite.internal.wire.Subscriber;
-import com.google.cloud.pubsublite.proto.Cursor;
+import com.google.cloud.pubsublite.internal.wire.SystemExecutors;
 import com.google.cloud.pubsublite.proto.FlowControlRequest;
-import com.google.cloud.pubsublite.proto.SeekRequest;
 import com.google.cloud.pubsublite.proto.SequencedMessage;
 import com.google.protobuf.util.Timestamps;
 import java.util.List;
@@ -36,19 +37,17 @@
 import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import java.util.function.Function;
-import org.apache.beam.sdk.io.range.OffsetRange;
 import org.apache.beam.sdk.transforms.DoFn.OutputReceiver;
 import org.apache.beam.sdk.transforms.DoFn.ProcessContinuation;
 import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.MoreExecutors;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.util.concurrent.SettableFuture;
 import org.joda.time.Duration;
 import org.joda.time.Instant;
 
 class SubscriptionPartitionProcessorImpl extends Listener
     implements SubscriptionPartitionProcessor {
-  private final RestrictionTracker<OffsetRange, OffsetByteProgress> tracker;
+  private final RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker;
   private final OutputReceiver<SequencedMessage> receiver;
   private final Subscriber subscriber;
   private final SettableFuture<Void> completionFuture = SettableFuture.create();
@@ -57,7 +56,7 @@
 
   @SuppressWarnings("methodref.receiver.bound.invalid")
   SubscriptionPartitionProcessorImpl(
-      RestrictionTracker<OffsetRange, OffsetByteProgress> tracker,
+      RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker,
       OutputReceiver<SequencedMessage> receiver,
       Function<Consumer<List<SequencedMessage>>, Subscriber> subscriberFactory,
       FlowControlSettings flowControlSettings) {
@@ -70,23 +69,15 @@
   @Override
   @SuppressWarnings("argument.type.incompatible")
   public void start() throws CheckedApiException {
-    this.subscriber.addListener(this, MoreExecutors.directExecutor());
+    this.subscriber.addListener(this, SystemExecutors.getFuturesExecutor());
     this.subscriber.startAsync();
     this.subscriber.awaitRunning();
     try {
-      this.subscriber
-          .seek(
-              SeekRequest.newBuilder()
-                  .setCursor(Cursor.newBuilder().setOffset(tracker.currentRestriction().getFrom()))
-                  .build())
-          .get();
       this.subscriber.allowFlow(
           FlowControlRequest.newBuilder()
               .setAllowedBytes(flowControlSettings.bytesOutstanding())
               .setAllowedMessages(flowControlSettings.messagesOutstanding())
               .build());
-    } catch (ExecutionException e) {
-      throw ExtractStatus.toCanonical(e.getCause());
     } catch (Throwable t) {
       throw ExtractStatus.toCanonical(t);
     }
@@ -125,7 +116,7 @@
 
   @Override
   public void close() {
-    subscriber.stopAsync().awaitTerminated();
+    blockingShutdown(subscriber);
   }
 
   @Override
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/TopicBacklogReaderSettings.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/TopicBacklogReaderSettings.java
index 8c1dd94..79db0f1 100644
--- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/TopicBacklogReaderSettings.java
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/TopicBacklogReaderSettings.java
@@ -62,7 +62,7 @@
       try (AdminClient adminClient =
           AdminClient.create(
               AdminClientSettings.newBuilder()
-                  .setRegion(subscriptionPath.location().region())
+                  .setRegion(subscriptionPath.location().extractRegion())
                   .build())) {
         return setTopicPath(
             TopicPath.parse(adminClient.getSubscription(subscriptionPath).get().getTopic()));
@@ -81,7 +81,9 @@
 
   TopicBacklogReader instantiate() throws ApiException {
     TopicStatsClientSettings settings =
-        TopicStatsClientSettings.newBuilder().setRegion(topicPath().location().region()).build();
+        TopicStatsClientSettings.newBuilder()
+            .setRegion(topicPath().location().extractRegion())
+            .build();
     TopicBacklogReader impl =
         new TopicBacklogReaderImpl(TopicStatsClient.create(settings), topicPath(), partition());
     return new LimitingTopicBacklogReader(impl, Ticker.systemTicker());
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/TrackerWithProgress.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/TrackerWithProgress.java
new file mode 100644
index 0000000..7f0d030
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/pubsublite/TrackerWithProgress.java
@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsublite;
+
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.HasProgress;
+
+public abstract class TrackerWithProgress
+    extends RestrictionTracker<OffsetByteRange, OffsetByteProgress> implements HasProgress {}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnPartitionQueryTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnPartitionQueryTest.java
index 0c9bbf1..1f29883 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnPartitionQueryTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/FirestoreV1FnPartitionQueryTest.java
@@ -99,6 +99,41 @@
     assertEquals(expected, allValues);
   }
 
+  @Test
+  public void endToEnd_emptyCursors() throws Exception {
+    // First page of the response
+    PartitionQueryRequest request1 =
+        PartitionQueryRequest.newBuilder()
+            .setParent(String.format("projects/%s/databases/(default)/document", projectId))
+            .build();
+    PartitionQueryResponse response1 = PartitionQueryResponse.newBuilder().build();
+    when(callable.call(request1)).thenReturn(pagedResponse1);
+    when(page1.getResponse()).thenReturn(response1);
+    when(pagedResponse1.iteratePages()).thenReturn(ImmutableList.of(page1));
+
+    when(stub.partitionQueryPagedCallable()).thenReturn(callable);
+
+    when(ff.getFirestoreStub(any())).thenReturn(stub);
+    RpcQosOptions options = RpcQosOptions.defaultOptions();
+    when(ff.getRpcQos(any()))
+        .thenReturn(FirestoreStatefulComponentFactory.INSTANCE.getRpcQos(options));
+
+    ArgumentCaptor<PartitionQueryPair> responses =
+        ArgumentCaptor.forClass(PartitionQueryPair.class);
+
+    doNothing().when(processContext).output(responses.capture());
+
+    when(processContext.element()).thenReturn(request1);
+
+    PartitionQueryFn fn = new PartitionQueryFn(clock, ff, options);
+
+    runFunction(fn);
+
+    List<PartitionQueryPair> expected = newArrayList(new PartitionQueryPair(request1, response1));
+    List<PartitionQueryPair> allValues = responses.getAllValues();
+    assertEquals(expected, allValues);
+  }
+
   @Override
   public void resumeFromLastReadValue() throws Exception {
     when(ff.getFirestoreStub(any())).thenReturn(stub);
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/PartitionQueryResponseToRunQueryRequestTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/PartitionQueryResponseToRunQueryRequestTest.java
index 25ed63c..ed789da 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/PartitionQueryResponseToRunQueryRequestTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/firestore/PartitionQueryResponseToRunQueryRequestTest.java
@@ -121,6 +121,39 @@
     assertEquals(expectedQueries, actualQueries);
   }
 
+  @Test
+  public void ensureCursorPairingWorks_emptyCursorsInResponse() {
+    StructuredQuery query =
+        StructuredQuery.newBuilder()
+            .addFrom(
+                CollectionSelector.newBuilder()
+                    .setAllDescendants(true)
+                    .setCollectionId("c1")
+                    .build())
+            .build();
+
+    List<StructuredQuery> expectedQueries = newArrayList(query);
+
+    PartitionQueryPair partitionQueryPair =
+        new PartitionQueryPair(
+            PartitionQueryRequest.newBuilder().setStructuredQuery(query).build(),
+            PartitionQueryResponse.newBuilder().build());
+
+    ArgumentCaptor<RunQueryRequest> captor = ArgumentCaptor.forClass(RunQueryRequest.class);
+    when(processContext.element()).thenReturn(partitionQueryPair);
+    doNothing().when(processContext).output(captor.capture());
+
+    PartitionQueryResponseToRunQueryRequest fn = new PartitionQueryResponseToRunQueryRequest();
+    fn.processElement(processContext);
+
+    List<StructuredQuery> actualQueries =
+        captor.getAllValues().stream()
+            .map(RunQueryRequest::getStructuredQuery)
+            .collect(Collectors.toList());
+
+    assertEquals(expectedQueries, actualQueries);
+  }
+
   private static Cursor referenceValueCursor(String referenceValue) {
     return Cursor.newBuilder()
         .addValues(Value.newBuilder().setReferenceValue(referenceValue).build())
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTrackerTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTrackerTest.java
index f34ebb6..5a31f4f 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTrackerTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/OffsetByteRangeTrackerTest.java
@@ -49,7 +49,7 @@
   private static final double IGNORED_FRACTION = -10000000.0;
   private static final long MIN_BYTES = 1000;
   private static final OffsetRange RANGE = new OffsetRange(123L, Long.MAX_VALUE);
-  private final TopicBacklogReader reader = mock(TopicBacklogReader.class);
+  private final TopicBacklogReader unownedBacklogReader = mock(TopicBacklogReader.class);
 
   @Spy Ticker ticker;
   private OffsetByteRangeTracker tracker;
@@ -60,14 +60,18 @@
     when(ticker.read()).thenReturn(0L);
     tracker =
         new OffsetByteRangeTracker(
-            RANGE, reader, Stopwatch.createUnstarted(ticker), Duration.millis(500), MIN_BYTES);
+            OffsetByteRange.of(RANGE, 0),
+            unownedBacklogReader,
+            Stopwatch.createUnstarted(ticker),
+            Duration.millis(500),
+            MIN_BYTES);
   }
 
   @Test
   public void progressTracked() {
     assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(123), 10)));
     assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(124), 11)));
-    when(reader.computeMessageStats(Offset.of(125)))
+    when(unownedBacklogReader.computeMessageStats(Offset.of(125)))
         .thenReturn(ComputeMessageStatsResponse.newBuilder().setMessageBytes(1000).build());
     Progress progress = tracker.getProgress();
     assertEquals(21, progress.getWorkCompleted(), .0001);
@@ -76,7 +80,7 @@
 
   @Test
   public void getProgressStatsFailure() {
-    when(reader.computeMessageStats(Offset.of(123)))
+    when(unownedBacklogReader.computeMessageStats(Offset.of(123)))
         .thenThrow(new CheckedApiException(Code.INTERNAL).underlying);
     assertThrows(ApiException.class, tracker::getProgress);
   }
@@ -86,11 +90,15 @@
   public void claimSplitSuccess() {
     assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(1_000), MIN_BYTES)));
     assertTrue(tracker.tryClaim(OffsetByteProgress.of(Offset.of(10_000), MIN_BYTES)));
-    SplitResult<OffsetRange> splits = tracker.trySplit(IGNORED_FRACTION);
-    assertEquals(RANGE.getFrom(), splits.getPrimary().getFrom());
-    assertEquals(10_001, splits.getPrimary().getTo());
-    assertEquals(10_001, splits.getResidual().getFrom());
-    assertEquals(Long.MAX_VALUE, splits.getResidual().getTo());
+    SplitResult<OffsetByteRange> splits = tracker.trySplit(IGNORED_FRACTION);
+    OffsetByteRange primary = splits.getPrimary();
+    assertEquals(RANGE.getFrom(), primary.getRange().getFrom());
+    assertEquals(10_001, primary.getRange().getTo());
+    assertEquals(MIN_BYTES * 2, primary.getByteCount());
+    OffsetByteRange residual = splits.getResidual();
+    assertEquals(10_001, residual.getRange().getFrom());
+    assertEquals(Long.MAX_VALUE, residual.getRange().getTo());
+    assertEquals(0, residual.getByteCount());
     assertEquals(splits.getPrimary(), tracker.currentRestriction());
     tracker.checkDone();
     assertNull(tracker.trySplit(IGNORED_FRACTION));
@@ -100,10 +108,10 @@
   @SuppressWarnings({"dereference.of.nullable", "argument.type.incompatible"})
   public void splitWithoutClaimEmpty() {
     when(ticker.read()).thenReturn(100000000000000L);
-    SplitResult<OffsetRange> splits = tracker.trySplit(IGNORED_FRACTION);
-    assertEquals(RANGE.getFrom(), splits.getPrimary().getFrom());
-    assertEquals(RANGE.getFrom(), splits.getPrimary().getTo());
-    assertEquals(RANGE, splits.getResidual());
+    SplitResult<OffsetByteRange> splits = tracker.trySplit(IGNORED_FRACTION);
+    assertEquals(RANGE.getFrom(), splits.getPrimary().getRange().getFrom());
+    assertEquals(RANGE.getFrom(), splits.getPrimary().getRange().getTo());
+    assertEquals(RANGE, splits.getResidual().getRange());
     assertEquals(splits.getPrimary(), tracker.currentRestriction());
     tracker.checkDone();
     assertNull(tracker.trySplit(IGNORED_FRACTION));
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdfTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdfTest.java
index 598037e..0a4e3e7 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdfTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/PerSubscriptionPartitionSdfTest.java
@@ -28,6 +28,7 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.initMocks;
 
@@ -51,6 +52,8 @@
 import org.apache.beam.sdk.transforms.SerializableBiFunction;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker.Progress;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.math.DoubleMath;
 import org.joda.time.Duration;
 import org.junit.Before;
 import org.junit.Test;
@@ -65,22 +68,24 @@
 public class PerSubscriptionPartitionSdfTest {
   private static final Duration MAX_SLEEP_TIME =
       Duration.standardMinutes(10).plus(Duration.millis(10));
-  private static final OffsetRange RESTRICTION = new OffsetRange(1, Long.MAX_VALUE);
+  private static final OffsetByteRange RESTRICTION =
+      OffsetByteRange.of(new OffsetRange(1, Long.MAX_VALUE), 0);
   private static final SubscriptionPartition PARTITION =
       SubscriptionPartition.of(example(SubscriptionPath.class), example(Partition.class));
 
   @Mock SerializableFunction<SubscriptionPartition, InitialOffsetReader> offsetReaderFactory;
 
+  @Mock ManagedBacklogReaderFactory backlogReaderFactory;
+  @Mock TopicBacklogReader backlogReader;
+
   @Mock
-  SerializableBiFunction<
-          SubscriptionPartition, OffsetRange, RestrictionTracker<OffsetRange, OffsetByteProgress>>
-      trackerFactory;
+  SerializableBiFunction<TopicBacklogReader, OffsetByteRange, TrackerWithProgress> trackerFactory;
 
   @Mock SubscriptionPartitionProcessorFactory processorFactory;
   @Mock SerializableFunction<SubscriptionPartition, Committer> committerFactory;
 
   @Mock InitialOffsetReader initialOffsetReader;
-  @Spy RestrictionTracker<OffsetRange, OffsetByteProgress> tracker;
+  @Spy TrackerWithProgress tracker;
   @Mock OutputReceiver<SequencedMessage> output;
   @Mock SubscriptionPartitionProcessor processor;
 
@@ -98,9 +103,11 @@
     when(trackerFactory.apply(any(), any())).thenReturn(tracker);
     when(committerFactory.apply(any())).thenReturn(committer);
     when(tracker.currentRestriction()).thenReturn(RESTRICTION);
+    when(backlogReaderFactory.newReader(any())).thenReturn(backlogReader);
     sdf =
         new PerSubscriptionPartitionSdf(
             MAX_SLEEP_TIME,
+            backlogReaderFactory,
             offsetReaderFactory,
             trackerFactory,
             processorFactory,
@@ -110,9 +117,10 @@
   @Test
   public void getInitialRestrictionReadSuccess() {
     when(initialOffsetReader.read()).thenReturn(example(Offset.class));
-    OffsetRange range = sdf.getInitialRestriction(PARTITION);
-    assertEquals(example(Offset.class).value(), range.getFrom());
-    assertEquals(Long.MAX_VALUE, range.getTo());
+    OffsetByteRange range = sdf.getInitialRestriction(PARTITION);
+    assertEquals(example(Offset.class).value(), range.getRange().getFrom());
+    assertEquals(Long.MAX_VALUE, range.getRange().getTo());
+    assertEquals(0, range.getByteCount());
     verify(offsetReaderFactory).apply(PARTITION);
   }
 
@@ -125,7 +133,13 @@
   @Test
   public void newTrackerCallsFactory() {
     assertSame(tracker, sdf.newTracker(PARTITION, RESTRICTION));
-    verify(trackerFactory).apply(PARTITION, RESTRICTION);
+    verify(trackerFactory).apply(backlogReader, RESTRICTION);
+  }
+
+  @Test
+  public void tearDownClosesBacklogReaderFactory() {
+    sdf.teardown();
+    verify(backlogReaderFactory).close();
   }
 
   @Test
@@ -159,12 +173,48 @@
     order2.verify(committer).awaitTerminated();
   }
 
+  private static final class NoopManagedBacklogReaderFactory
+      implements ManagedBacklogReaderFactory {
+    @Override
+    public TopicBacklogReader newReader(SubscriptionPartition subscriptionPartition) {
+      return null;
+    }
+
+    @Override
+    public void close() {}
+  }
+
   @Test
   @SuppressWarnings("return.type.incompatible")
   public void dofnIsSerializable() throws Exception {
     ObjectOutputStream output = new ObjectOutputStream(new ByteArrayOutputStream());
     output.writeObject(
         new PerSubscriptionPartitionSdf(
-            MAX_SLEEP_TIME, x -> null, (x, y) -> null, (x, y, z) -> null, (x) -> null));
+            MAX_SLEEP_TIME,
+            new NoopManagedBacklogReaderFactory(),
+            x -> null,
+            (x, y) -> null,
+            (x, y, z) -> null,
+            (x) -> null));
+  }
+
+  @Test
+  public void getProgressUnboundedRangeDelegates() {
+    Progress progress = Progress.from(0, 0.2);
+    when(tracker.getProgress()).thenReturn(progress);
+    assertTrue(
+        DoubleMath.fuzzyEquals(
+            progress.getWorkRemaining(), sdf.getSize(PARTITION, RESTRICTION), .0001));
+    verify(tracker).getProgress();
+  }
+
+  @Test
+  public void getProgressBoundedReturnsBytes() {
+    assertTrue(
+        DoubleMath.fuzzyEquals(
+            123.0,
+            sdf.getSize(PARTITION, OffsetByteRange.of(new OffsetRange(87, 8000), 123)),
+            .0001));
+    verifyNoInteractions(tracker);
   }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/ReadWriteIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/ReadWriteIT.java
new file mode 100644
index 0000000..e242942
--- /dev/null
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/ReadWriteIT.java
@@ -0,0 +1,280 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.gcp.pubsublite;
+
+import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+import static org.junit.Assert.fail;
+
+import com.google.cloud.pubsublite.AdminClient;
+import com.google.cloud.pubsublite.AdminClientSettings;
+import com.google.cloud.pubsublite.BacklogLocation;
+import com.google.cloud.pubsublite.CloudZone;
+import com.google.cloud.pubsublite.Message;
+import com.google.cloud.pubsublite.ProjectId;
+import com.google.cloud.pubsublite.SubscriptionName;
+import com.google.cloud.pubsublite.SubscriptionPath;
+import com.google.cloud.pubsublite.TopicName;
+import com.google.cloud.pubsublite.TopicPath;
+import com.google.cloud.pubsublite.proto.PubSubMessage;
+import com.google.cloud.pubsublite.proto.SequencedMessage;
+import com.google.cloud.pubsublite.proto.Subscription;
+import com.google.cloud.pubsublite.proto.Subscription.DeliveryConfig.DeliveryRequirement;
+import com.google.cloud.pubsublite.proto.Topic;
+import com.google.cloud.pubsublite.proto.Topic.PartitionConfig.Capacity;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import com.google.protobuf.ByteString;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.StreamingOptions;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.FlatMapElements;
+import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.SimpleFunction;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
+import org.joda.time.Duration;
+import org.junit.After;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@RunWith(JUnit4.class)
+public class ReadWriteIT {
+  private static final Logger LOG = LoggerFactory.getLogger(ReadWriteIT.class);
+  private static final CloudZone ZONE = CloudZone.parse("us-central1-b");
+  private static final int MESSAGE_COUNT = 90;
+
+  @Rule public transient TestPipeline pipeline = TestPipeline.create();
+
+  private static ProjectId getProject(PipelineOptions options) {
+    return ProjectId.of(checkArgumentNotNull(options.as(GcpOptions.class).getProject()));
+  }
+
+  private static String randomName() {
+    return "beam_it_resource_" + ThreadLocalRandom.current().nextLong();
+  }
+
+  private static AdminClient newAdminClient() {
+    return AdminClient.create(AdminClientSettings.newBuilder().setRegion(ZONE.region()).build());
+  }
+
+  private final Deque<Runnable> cleanupActions = new ArrayDeque<>();
+
+  private TopicPath createTopic(ProjectId id) throws Exception {
+    TopicPath toReturn =
+        TopicPath.newBuilder()
+            .setProject(id)
+            .setLocation(ZONE)
+            .setName(TopicName.of(randomName()))
+            .build();
+    Topic.Builder topic = Topic.newBuilder().setName(toReturn.toString());
+    topic
+        .getPartitionConfigBuilder()
+        .setCount(2)
+        .setCapacity(Capacity.newBuilder().setPublishMibPerSec(4).setSubscribeMibPerSec(4));
+    topic.getRetentionConfigBuilder().setPerPartitionBytes(30 * (1L << 30));
+    cleanupActions.addLast(
+        () -> {
+          try (AdminClient client = newAdminClient()) {
+            client.deleteTopic(toReturn).get();
+          } catch (Throwable t) {
+            LOG.error("Failed to clean up topic.", t);
+          }
+        });
+    try (AdminClient client = newAdminClient()) {
+      client.createTopic(topic.build()).get();
+    }
+    return toReturn;
+  }
+
+  private SubscriptionPath createSubscription(TopicPath topic) throws Exception {
+    SubscriptionPath toReturn =
+        SubscriptionPath.newBuilder()
+            .setProject(topic.project())
+            .setLocation(ZONE)
+            .setName(SubscriptionName.of(randomName()))
+            .build();
+    Subscription.Builder subscription = Subscription.newBuilder().setName(toReturn.toString());
+    subscription
+        .getDeliveryConfigBuilder()
+        .setDeliveryRequirement(DeliveryRequirement.DELIVER_IMMEDIATELY);
+    subscription.setTopic(topic.toString());
+    cleanupActions.addLast(
+        () -> {
+          try (AdminClient client = newAdminClient()) {
+            client.deleteSubscription(toReturn).get();
+          } catch (Throwable t) {
+            LOG.error("Failed to clean up subscription.", t);
+          }
+        });
+    try (AdminClient client = newAdminClient()) {
+      client.createSubscription(subscription.build(), BacklogLocation.BEGINNING).get();
+    }
+    return toReturn;
+  }
+
+  @After
+  public void tearDown() {
+    while (!cleanupActions.isEmpty()) {
+      cleanupActions.removeLast().run();
+    }
+  }
+
+  // Workaround for BEAM-12867
+  // TODO(BEAM-12867): Remove this.
+  private static class CustomCreate extends PTransform<PCollection<Void>, PCollection<Integer>> {
+    @Override
+    public PCollection<Integer> expand(PCollection<Void> input) {
+      return input.apply(
+          "createIndexes",
+          FlatMapElements.via(
+              new SimpleFunction<Void, Iterable<Integer>>() {
+                @Override
+                public Iterable<Integer> apply(Void input) {
+                  return IntStream.range(0, MESSAGE_COUNT).boxed().collect(Collectors.toList());
+                }
+              }));
+    }
+  }
+
+  public static void writeMessages(TopicPath topicPath, Pipeline pipeline) {
+    PCollection<Void> trigger = pipeline.apply(Create.of((Void) null));
+    PCollection<Integer> indexes = trigger.apply("createIndexes", new CustomCreate());
+    PCollection<PubSubMessage> messages =
+        indexes.apply(
+            "createMessages",
+            MapElements.via(
+                new SimpleFunction<Integer, PubSubMessage>(
+                    index ->
+                        Message.builder()
+                            .setData(ByteString.copyFromUtf8(index.toString()))
+                            .build()
+                            .toProto()) {}));
+    // Add UUIDs to messages for later deduplication.
+    messages = messages.apply("addUuids", PubsubLiteIO.addUuids());
+    messages.apply(
+        "writeMessages",
+        PubsubLiteIO.write(PublisherOptions.newBuilder().setTopicPath(topicPath).build()));
+  }
+
+  public static PCollection<SequencedMessage> readMessages(
+      SubscriptionPath subscriptionPath, Pipeline pipeline) {
+    PCollection<SequencedMessage> messages =
+        pipeline.apply(
+            "readMessages",
+            PubsubLiteIO.read(
+                SubscriberOptions.newBuilder()
+                    .setSubscriptionPath(subscriptionPath)
+                    // setMinBundleTimeout INTENDED FOR TESTING ONLY
+                    // This sacrifices efficiency to make tests run faster. Do not use this in a
+                    // real pipeline!
+                    .setMinBundleTimeout(Duration.standardSeconds(5))
+                    .build()));
+    // Deduplicate messages based on the uuids added in PubsubLiteIO.addUuids() when writing.
+    return messages.apply(
+        "dedupeMessages", PubsubLiteIO.deduplicate(UuidDeduplicationOptions.newBuilder().build()));
+  }
+
+  // This static out of band communication is needed to retain serializability.
+  @GuardedBy("ReadWriteIT.class")
+  private static final List<SequencedMessage> received = new ArrayList<>();
+
+  private static synchronized void addMessageReceived(SequencedMessage message) {
+    received.add(message);
+  }
+
+  private static synchronized List<SequencedMessage> getTestQuickstartReceived() {
+    return ImmutableList.copyOf(received);
+  }
+
+  private static PTransform<PCollection<? extends SequencedMessage>, PCollection<Void>>
+      collectTestQuickstart() {
+    return MapElements.via(
+        new SimpleFunction<SequencedMessage, Void>() {
+          @Override
+          public Void apply(SequencedMessage input) {
+            addMessageReceived(input);
+            return null;
+          }
+        });
+  }
+
+  @Test
+  public void testReadWrite() throws Exception {
+    pipeline.getOptions().as(StreamingOptions.class).setStreaming(true);
+    pipeline.getOptions().as(TestPipelineOptions.class).setBlockOnRun(false);
+
+    TopicPath topic = createTopic(getProject(pipeline.getOptions()));
+    SubscriptionPath subscription = createSubscription(topic);
+
+    // Publish some messages
+    writeMessages(topic, pipeline);
+
+    // Read some messages. They should be deduplicated by the time we see them, so there should be
+    // exactly numMessages, one for every index in [0,MESSAGE_COUNT).
+    PCollection<SequencedMessage> messages = readMessages(subscription, pipeline);
+    messages.apply("messageReceiver", collectTestQuickstart());
+    pipeline.run();
+    LOG.info("Running!");
+    for (int round = 0; round < 120; ++round) {
+      Thread.sleep(1000);
+      Map<Integer, Integer> receivedCounts = new HashMap<>();
+      for (SequencedMessage message : getTestQuickstartReceived()) {
+        int id = Integer.parseInt(message.getMessage().getData().toStringUtf8());
+        receivedCounts.put(id, receivedCounts.getOrDefault(id, 0) + 1);
+      }
+      LOG.info("Performing comparison round {}.\n", round);
+      boolean done = true;
+      List<Integer> missing = new ArrayList<>();
+      for (int id = 0; id < MESSAGE_COUNT; id++) {
+        int idCount = receivedCounts.getOrDefault(id, 0);
+        if (idCount == 0) {
+          missing.add(id);
+          done = false;
+        }
+        if (idCount > 1) {
+          fail(String.format("Failed to deduplicate message with id %s.", id));
+        }
+      }
+      LOG.info("Still messing messages: {}.\n", missing);
+      if (done) {
+        return;
+      }
+    }
+    fail(
+        String.format(
+            "Failed to receive all messages after 2 minutes. Received %s messages.",
+            getTestQuickstartReceived().size()));
+  }
+}
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImplTest.java
index dbf3b93..3d743758 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImplTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/pubsublite/SubscriptionPartitionProcessorImplTest.java
@@ -31,7 +31,6 @@
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.initMocks;
 
-import com.google.api.core.ApiFutures;
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.cloud.pubsublite.Offset;
@@ -40,7 +39,6 @@
 import com.google.cloud.pubsublite.internal.wire.Subscriber;
 import com.google.cloud.pubsublite.proto.Cursor;
 import com.google.cloud.pubsublite.proto.FlowControlRequest;
-import com.google.cloud.pubsublite.proto.SeekRequest;
 import com.google.cloud.pubsublite.proto.SequencedMessage;
 import com.google.protobuf.util.Timestamps;
 import java.util.List;
@@ -64,7 +62,7 @@
 @RunWith(JUnit4.class)
 @SuppressWarnings("initialization.fields.uninitialized")
 public class SubscriptionPartitionProcessorImplTest {
-  @Spy RestrictionTracker<OffsetRange, OffsetByteProgress> tracker;
+  @Spy RestrictionTracker<OffsetByteRange, OffsetByteProgress> tracker;
   @Mock OutputReceiver<SequencedMessage> receiver;
   @Mock Function<Consumer<List<SequencedMessage>>, Subscriber> subscriberFactory;
 
@@ -83,6 +81,10 @@
         .build();
   }
 
+  private OffsetByteRange initialRange() {
+    return OffsetByteRange.of(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE));
+  }
+
   @Before
   public void setUp() {
     initMocks(this);
@@ -100,18 +102,11 @@
 
   @Test
   public void lifecycle() throws Exception {
-    when(tracker.currentRestriction())
-        .thenReturn(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE));
-    when(subscriber.seek(any())).thenReturn(ApiFutures.immediateFuture(example(Offset.class)));
+    when(tracker.currentRestriction()).thenReturn(initialRange());
     processor.start();
     verify(subscriber).startAsync();
     verify(subscriber).awaitRunning();
     verify(subscriber)
-        .seek(
-            SeekRequest.newBuilder()
-                .setCursor(Cursor.newBuilder().setOffset(example(Offset.class).value()))
-                .build());
-    verify(subscriber)
         .allowFlow(
             FlowControlRequest.newBuilder()
                 .setAllowedBytes(DEFAULT_FLOW_CONTROL.bytesOutstanding())
@@ -123,29 +118,15 @@
   }
 
   @Test
-  public void lifecycleSeekThrows() throws Exception {
-    when(tracker.currentRestriction())
-        .thenReturn(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE));
-    when(subscriber.seek(any()))
-        .thenReturn(ApiFutures.immediateFailedFuture(new CheckedApiException(Code.OUT_OF_RANGE)));
+  public void lifecycleFlowControlThrows() throws Exception {
+    when(tracker.currentRestriction()).thenReturn(initialRange());
     doThrow(new CheckedApiException(Code.OUT_OF_RANGE)).when(subscriber).allowFlow(any());
     assertThrows(CheckedApiException.class, () -> processor.start());
   }
 
   @Test
-  public void lifecycleFlowControlThrows() {
-    when(tracker.currentRestriction())
-        .thenReturn(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE));
-    when(subscriber.seek(any()))
-        .thenReturn(ApiFutures.immediateFailedFuture(new CheckedApiException(Code.OUT_OF_RANGE)));
-    assertThrows(CheckedApiException.class, () -> processor.start());
-  }
-
-  @Test
   public void lifecycleSubscriberAwaitThrows() throws Exception {
-    when(tracker.currentRestriction())
-        .thenReturn(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE));
-    when(subscriber.seek(any())).thenReturn(ApiFutures.immediateFuture(example(Offset.class)));
+    when(tracker.currentRestriction()).thenReturn(initialRange());
     processor.start();
     doThrow(new CheckedApiException(Code.INTERNAL).underlying).when(subscriber).awaitTerminated();
     assertThrows(ApiException.class, () -> processor.close());
@@ -155,21 +136,19 @@
 
   @Test
   public void subscriberFailureFails() throws Exception {
-    when(tracker.currentRestriction())
-        .thenReturn(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE));
-    when(subscriber.seek(any())).thenReturn(ApiFutures.immediateFuture(example(Offset.class)));
+    when(tracker.currentRestriction()).thenReturn(initialRange());
     processor.start();
     subscriber.fail(new CheckedApiException(Code.OUT_OF_RANGE));
     ApiException e =
-        assertThrows(ApiException.class, () -> processor.waitForCompletion(Duration.ZERO));
+        assertThrows(
+            // Longer wait is needed due to listener asynchrony.
+            ApiException.class, () -> processor.waitForCompletion(Duration.standardSeconds(1)));
     assertEquals(Code.OUT_OF_RANGE, e.getStatusCode().getCode());
   }
 
   @Test
   public void allowFlowFailureFails() throws Exception {
-    when(tracker.currentRestriction())
-        .thenReturn(new OffsetRange(example(Offset.class).value(), Long.MAX_VALUE));
-    when(subscriber.seek(any())).thenReturn(ApiFutures.immediateFuture(example(Offset.class)));
+    when(tracker.currentRestriction()).thenReturn(initialRange());
     processor.start();
     when(tracker.tryClaim(any())).thenReturn(true);
     doThrow(new CheckedApiException(Code.OUT_OF_RANGE)).when(subscriber).allowFlow(any());
diff --git a/sdks/java/io/jms/build.gradle b/sdks/java/io/jms/build.gradle
index 69cd88f..3f9d0c7 100644
--- a/sdks/java/io/jms/build.gradle
+++ b/sdks/java/io/jms/build.gradle
@@ -36,6 +36,7 @@
   testCompile library.java.activemq_kahadb_store
   testCompile library.java.activemq_client
   testCompile library.java.junit
+  testCompile library.java.mockito_core
   testRuntimeOnly library.java.slf4j_jdk14
   testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")
 }
diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/AutoScaler.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/AutoScaler.java
new file mode 100644
index 0000000..0e023d1
--- /dev/null
+++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/AutoScaler.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.jms;
+
+import java.io.Serializable;
+import org.apache.beam.sdk.io.UnboundedSource;
+
+/**
+ * Enables users to specify their own `JMS` backlog reporters enabling {@link JmsIO} to report
+ * {@link UnboundedSource.UnboundedReader#getTotalBacklogBytes()}.
+ */
+public interface AutoScaler extends Serializable {
+
+  /** The {@link AutoScaler} is started when the {@link JmsIO.UnboundedJmsReader} is started. */
+  void start();
+
+  /**
+   * Returns the size of the backlog of unread data in the underlying data source represented by all
+   * splits of this source.
+   */
+  long getTotalBacklogBytes();
+
+  /** The {@link AutoScaler} is stopped when the {@link JmsIO.UnboundedJmsReader} is closed. */
+  void stop();
+}
diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/DefaultAutoscaler.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/DefaultAutoscaler.java
new file mode 100644
index 0000000..2b05cf6
--- /dev/null
+++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/DefaultAutoscaler.java
@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.sdk.io.jms;
+
+import static org.apache.beam.sdk.io.UnboundedSource.UnboundedReader.BACKLOG_UNKNOWN;
+
+/**
+ * Default implementation of {@link AutoScaler}. Returns {@link
+ * org.apache.beam.sdk.io.UnboundedSource.UnboundedReader#BACKLOG_UNKNOWN} as the default value.
+ */
+public class DefaultAutoscaler implements AutoScaler {
+  @Override
+  public void start() {}
+
+  @Override
+  public long getTotalBacklogBytes() {
+    return BACKLOG_UNKNOWN;
+  }
+
+  @Override
+  public void stop() {}
+}
diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
index 4999e10..9fa4492 100644
--- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
+++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
@@ -196,6 +196,8 @@
 
     abstract @Nullable Coder<T> getCoder();
 
+    abstract @Nullable AutoScaler getAutoScaler();
+
     abstract Builder<T> builder();
 
     @AutoValue.Builder
@@ -218,6 +220,8 @@
 
       abstract Builder<T> setCoder(Coder<T> coder);
 
+      abstract Builder<T> setAutoScaler(AutoScaler autoScaler);
+
       abstract Read<T> build();
     }
 
@@ -344,6 +348,14 @@
       return builder().setCoder(coder).build();
     }
 
+    /**
+     * Sets the {@link AutoScaler} to use for reporting backlog during the execution of this source.
+     */
+    public Read<T> withAutoScaler(AutoScaler autoScaler) {
+      checkArgument(autoScaler != null, "autoScaler can not be null");
+      return builder().setAutoScaler(autoScaler).build();
+    }
+
     @Override
     public PCollection<T> expand(PBegin input) {
       checkArgument(getConnectionFactory() != null, "withConnectionFactory() is required");
@@ -447,6 +459,7 @@
     private Connection connection;
     private Session session;
     private MessageConsumer consumer;
+    private AutoScaler autoScaler;
 
     private T currentMessage;
     private Instant currentTimestamp;
@@ -474,6 +487,12 @@
         }
         connection.start();
         this.connection = connection;
+        if (spec.getAutoScaler() == null) {
+          this.autoScaler = new DefaultAutoscaler();
+        } else {
+          this.autoScaler = spec.getAutoScaler();
+        }
+        this.autoScaler.start();
       } catch (Exception e) {
         throw new IOException("Error connecting to JMS", e);
       }
@@ -545,6 +564,11 @@
     }
 
     @Override
+    public long getTotalBacklogBytes() {
+      return this.autoScaler.getTotalBacklogBytes();
+    }
+
+    @Override
     public UnboundedSource<T, ?> getCurrentSource() {
       return source;
     }
@@ -565,6 +589,10 @@
           connection.close();
           connection = null;
         }
+        if (autoScaler != null) {
+          autoScaler.stop();
+          autoScaler = null;
+        }
       } catch (Exception e) {
         throw new IOException(e);
       }
diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java
index c335f8a..a9f3c3f 100644
--- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java
+++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java
@@ -17,12 +17,17 @@
  */
 package org.apache.beam.sdk.io.jms;
 
+import static org.apache.beam.sdk.io.UnboundedSource.UnboundedReader.BACKLOG_UNKNOWN;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import java.io.IOException;
 import java.lang.reflect.Proxy;
@@ -421,6 +426,50 @@
     CoderProperties.coderDecodeEncodeEqual(coder, jmsCheckpointMark);
   }
 
+  @Test
+  public void testDefaultAutoscaler() throws IOException {
+    JmsIO.Read spec =
+        JmsIO.read()
+            .withConnectionFactory(connectionFactory)
+            .withUsername(USERNAME)
+            .withPassword(PASSWORD)
+            .withQueue(QUEUE);
+    JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec);
+    JmsIO.UnboundedJmsReader reader = source.createReader(null, null);
+
+    // start the reader and check getSplitBacklogBytes and getTotalBacklogBytes values
+    reader.start();
+    assertEquals(BACKLOG_UNKNOWN, reader.getSplitBacklogBytes());
+    assertEquals(BACKLOG_UNKNOWN, reader.getTotalBacklogBytes());
+    reader.close();
+  }
+
+  @Test
+  public void testCustomAutoscaler() throws IOException {
+    long excpectedTotalBacklogBytes = 1111L;
+
+    AutoScaler autoScaler = mock(DefaultAutoscaler.class);
+    when(autoScaler.getTotalBacklogBytes()).thenReturn(excpectedTotalBacklogBytes);
+    JmsIO.Read spec =
+        JmsIO.read()
+            .withConnectionFactory(connectionFactory)
+            .withUsername(USERNAME)
+            .withPassword(PASSWORD)
+            .withQueue(QUEUE)
+            .withAutoScaler(autoScaler);
+
+    JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec);
+    JmsIO.UnboundedJmsReader reader = source.createReader(null, null);
+
+    // start the reader and check getSplitBacklogBytes and getTotalBacklogBytes values
+    reader.start();
+    verify(autoScaler, times(1)).start();
+    assertEquals(excpectedTotalBacklogBytes, reader.getTotalBacklogBytes());
+    verify(autoScaler, times(1)).getTotalBacklogBytes();
+    reader.close();
+    verify(autoScaler, times(1)).stop();
+  }
+
   private int count(String queue) throws Exception {
     Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD);
     connection.start();
diff --git a/sdks/python/apache_beam/runners/portability/expansion_service_test.py b/sdks/python/apache_beam/runners/portability/expansion_service_test.py
index 1c02f71..98d2faa 100644
--- a/sdks/python/apache_beam/runners/portability/expansion_service_test.py
+++ b/sdks/python/apache_beam/runners/portability/expansion_service_test.py
@@ -185,9 +185,12 @@
 @ptransform.PTransform.register_urn(TEST_COMPK_URN, None)
 class CombinePerKeyTransform(ptransform.PTransform):
   def expand(self, pcoll):
-    return pcoll \
-           | beam.CombinePerKey(sum).with_output_types(
-               typing.Tuple[str, int])
+    output = pcoll \
+           | beam.CombinePerKey(sum)
+    # TODO: Use `with_output_types` instead of explicitly
+    #  assigning to `.element_type` after fixing BEAM-12872
+    output.element_type = beam.typehints.Tuple[str, int]
+    return output
 
   def to_runner_api_parameter(self, unused_context):
     return TEST_COMPK_URN, None
diff --git a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py
index bebde98..9b0d6ff 100644
--- a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py
+++ b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server.py
@@ -63,12 +63,16 @@
       if not os.path.exists(self._executable_jar):
         parsed = urllib.parse.urlparse(self._executable_jar)
         if not parsed.scheme:
+          try:
+            flink_version = self.flink_version()
+          except Exception:
+            flink_version = '$FLINK_VERSION'
           raise ValueError(
               'Unable to parse jar URL "%s". If using a full URL, make sure '
               'the scheme is specified. If using a local file path, make sure '
               'the file exists; you may have to first build the job server '
               'using `./gradlew runners:flink:%s:job-server:shadowJar`.' %
-              (self._executable_jar, self._flink_version))
+              (self._executable_jar, flink_version))
       url = self._executable_jar
     else:
       url = job_server.JavaJarJobServer.path_to_beam_jar(
diff --git a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py
index 281e2d9..1294f46 100644
--- a/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py
+++ b/sdks/python/apache_beam/runners/portability/flink_uber_jar_job_server_test.py
@@ -192,6 +192,36 @@
     self.assertEqual(
         options_proto['beam:option:unknown_option_foo:v1'], 'some_value')
 
+  @requests_mock.mock()
+  def test_bad_url_flink_version(self, http_mock):
+    http_mock.get('http://flink/v1/config', json={'flink-version': '1.2.3.4'})
+    options = pipeline_options.FlinkRunnerOptions()
+    options.flink_job_server_jar = "bad url"
+    job_server = flink_uber_jar_job_server.FlinkUberJarJobServer(
+        'http://flink', options)
+    with self.assertRaises(ValueError) as context:
+      job_server.executable_jar()
+    self.assertEqual(
+        'Unable to parse jar URL "bad url". If using a full URL, make sure '
+        'the scheme is specified. If using a local file path, make sure '
+        'the file exists; you may have to first build the job server '
+        'using `./gradlew runners:flink:1.2:job-server:shadowJar`.',
+        str(context.exception))
+
+  def test_bad_url_placeholder_version(self):
+    options = pipeline_options.FlinkRunnerOptions()
+    options.flink_job_server_jar = "bad url"
+    job_server = flink_uber_jar_job_server.FlinkUberJarJobServer(
+        'http://example.com/bad', options)
+    with self.assertRaises(ValueError) as context:
+      job_server.executable_jar()
+    self.assertEqual(
+        'Unable to parse jar URL "bad url". If using a full URL, make sure '
+        'the scheme is specified. If using a local file path, make sure '
+        'the file exists; you may have to first build the job server '
+        'using `./gradlew runners:flink:$FLINK_VERSION:job-server:shadowJar`.',
+        str(context.exception))
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)
diff --git a/sdks/python/apache_beam/transforms/combiners.py b/sdks/python/apache_beam/transforms/combiners.py
index bcedd86..41ad3df 100644
--- a/sdks/python/apache_beam/transforms/combiners.py
+++ b/sdks/python/apache_beam/transforms/combiners.py
@@ -21,6 +21,7 @@
 
 import copy
 import heapq
+import itertools
 import operator
 import random
 from typing import Any
@@ -168,7 +169,8 @@
   """Combiners for obtaining extremal elements."""
 
   # pylint: disable=no-self-argument
-
+  @with_input_types(T)
+  @with_output_types(List[T])
   class Of(CombinerWithoutDefaults):
     """Obtain a list of the compare-most N elements in a PCollection.
 
@@ -202,10 +204,12 @@
         # This is a more efficient global algorithm.
         top_per_bundle = pcoll | core.ParDo(
             _TopPerBundle(self._n, self._key, self._reverse))
-        # If pcoll is empty, we can't guerentee that top_per_bundle
+        # If pcoll is empty, we can't guarantee that top_per_bundle
         # won't be empty, so inject at least one empty accumulator
-        # so that downstream is guerenteed to produce non-empty output.
-        empty_bundle = pcoll.pipeline | core.Create([(None, [])])
+        # so that downstream is guaranteed to produce non-empty output.
+        empty_bundle = (
+            pcoll.pipeline | core.Create([(None, [])]).with_output_types(
+                top_per_bundle.element_type))
         return ((top_per_bundle, empty_bundle) | core.Flatten()
                 | core.GroupByKey()
                 | core.ParDo(
@@ -219,6 +223,8 @@
               TopCombineFn(self._n, self._key,
                            self._reverse)).without_defaults()
 
+  @with_input_types(Tuple[K, V])
+  @with_output_types(Tuple[K, List[V]])
   class PerKey(ptransform.PTransform):
     """Identifies the compare-most N elements associated with each key.
 
@@ -524,6 +530,8 @@
 
   # pylint: disable=no-self-argument
 
+  @with_input_types(T)
+  @with_output_types(List[T])
   class FixedSizeGlobally(CombinerWithoutDefaults):
     """Sample n elements from the input PCollection without replacement."""
     def __init__(self, n):
@@ -543,6 +551,8 @@
     def default_label(self):
       return 'FixedSizeGlobally(%d)' % self._n
 
+  @with_input_types(Tuple[K, V])
+  @with_output_types(Tuple[K, List[V]])
   class FixedSizePerKey(ptransform.PTransform):
     """Sample n elements associated with each key without replacement."""
     def __init__(self, n):
@@ -597,16 +607,25 @@
 
 
 class _TupleCombineFnBase(core.CombineFn):
-  def __init__(self, *combiners):
+  def __init__(self, *combiners, merge_accumulators_batch_size=None):
     self._combiners = [core.CombineFn.maybe_from_callable(c) for c in combiners]
     self._named_combiners = combiners
+    # If the `merge_accumulators_batch_size` value is not specified, we chose a
+    # bounded default that is inversely proportional to the number of
+    # accumulators in merged tuples.
+    num_combiners = max(1, len(combiners))
+    self._merge_accumulators_batch_size = (
+        merge_accumulators_batch_size or max(10, 1000 // num_combiners))
 
   def display_data(self):
     combiners = [
         c.__name__ if hasattr(c, '__name__') else c.__class__.__name__
         for c in self._named_combiners
     ]
-    return {'combiners': str(combiners)}
+    return {
+        'combiners': str(combiners),
+        'merge_accumulators_batch_size': self._merge_accumulators_batch_size
+    }
 
   def setup(self, *args, **kwargs):
     for c in self._combiners:
@@ -616,10 +635,22 @@
     return [c.create_accumulator(*args, **kwargs) for c in self._combiners]
 
   def merge_accumulators(self, accumulators, *args, **kwargs):
-    return [
-        c.merge_accumulators(a, *args, **kwargs) for c,
-        a in zip(self._combiners, zip(*accumulators))
-    ]
+    # Make sure that `accumulators` is an iterator (so that the position is
+    # remembered).
+    accumulators = iter(accumulators)
+    result = next(accumulators)
+    while True:
+      # Load accumulators into memory and merge in batches to decrease peak
+      # memory usage.
+      accumulators_batch = [result] + list(
+          itertools.islice(accumulators, self._merge_accumulators_batch_size))
+      if len(accumulators_batch) == 1:
+        break
+      result = [
+          c.merge_accumulators(a, *args, **kwargs) for c,
+          a in zip(self._combiners, zip(*accumulators_batch))
+      ]
+    return result
 
   def compact(self, accumulator, *args, **kwargs):
     return [
@@ -670,6 +701,8 @@
     ]
 
 
+@with_input_types(T)
+@with_output_types(List[T])
 class ToList(CombinerWithoutDefaults):
   """A global CombineFn that condenses a PCollection into a single list."""
   def expand(self, pcoll):
@@ -698,6 +731,8 @@
     return accumulator
 
 
+@with_input_types(Tuple[K, V])
+@with_output_types(Dict[K, V])
 class ToDict(CombinerWithoutDefaults):
   """A global CombineFn that condenses a PCollection into a single dict.
 
@@ -735,6 +770,8 @@
     return accumulator
 
 
+@with_input_types(T)
+@with_output_types(Set[T])
 class ToSet(CombinerWithoutDefaults):
   """A global CombineFn that condenses a PCollection into a set."""
   def expand(self, pcoll):
diff --git a/sdks/python/apache_beam/transforms/combiners_test.py b/sdks/python/apache_beam/transforms/combiners_test.py
index d826287..7e0e835 100644
--- a/sdks/python/apache_beam/transforms/combiners_test.py
+++ b/sdks/python/apache_beam/transforms/combiners_test.py
@@ -249,7 +249,8 @@
     dd = DisplayData.create_from(transform)
     expected_items = [
         DisplayDataItemMatcher('combine_fn', combine.TupleCombineFn),
-        DisplayDataItemMatcher('combiners', "['max', 'MeanCombineFn', 'sum']")
+        DisplayDataItemMatcher('combiners', "['max', 'MeanCombineFn', 'sum']"),
+        DisplayDataItemMatcher('merge_accumulators_batch_size', 333),
     ]
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))
 
@@ -358,6 +359,49 @@
                   max).with_common_input()).without_defaults())
       assert_that(result, equal_to([(1, 7.0 / 4, 3)]))
 
+  def test_empty_tuple_combine_fn(self):
+    with TestPipeline() as p:
+      result = (
+          p
+          | Create([(), (), ()])
+          | beam.CombineGlobally(combine.TupleCombineFn()))
+      assert_that(result, equal_to([()]))
+
+  def test_tuple_combine_fn_batched_merge(self):
+    num_combine_fns = 10
+    max_num_accumulators_in_memory = 30
+    # Maximum number of accumulator tuples in memory - 1 for the merge result.
+    merge_accumulators_batch_size = (
+        max_num_accumulators_in_memory // num_combine_fns - 1)
+    num_accumulator_tuples_to_merge = 20
+
+    class CountedAccumulator:
+      count = 0
+      oom = False
+
+      def __init__(self):
+        if CountedAccumulator.count > max_num_accumulators_in_memory:
+          CountedAccumulator.oom = True
+        else:
+          CountedAccumulator.count += 1
+
+    class CountedAccumulatorCombineFn(beam.CombineFn):
+      def create_accumulator(self):
+        return CountedAccumulator()
+
+      def merge_accumulators(self, accumulators):
+        CountedAccumulator.count += 1
+        for _ in accumulators:
+          CountedAccumulator.count -= 1
+
+    combine_fn = combine.TupleCombineFn(
+        *[CountedAccumulatorCombineFn() for _ in range(num_combine_fns)],
+        merge_accumulators_batch_size=merge_accumulators_batch_size)
+    combine_fn.merge_accumulators(
+        combine_fn.create_accumulator()
+        for _ in range(num_accumulator_tuples_to_merge))
+    assert not CountedAccumulator.oom
+
   def test_to_list_and_to_dict1(self):
     with TestPipeline() as pipeline:
       the_list = [6, 3, 1, 1, 9, 1, 5, 2, 0, 6]
diff --git a/sdks/python/apache_beam/transforms/ptransform_test.py b/sdks/python/apache_beam/transforms/ptransform_test.py
index ec001e2..9ce33f5 100644
--- a/sdks/python/apache_beam/transforms/ptransform_test.py
+++ b/sdks/python/apache_beam/transforms/ptransform_test.py
@@ -2333,9 +2333,8 @@
 
     self.assertStartswith(
         e.exception.args[0],
-        "Type hint violation for 'CombinePerKey(TopCombineFn)': "
-        "requires Tuple[TypeVariable[K], TypeVariable[T]] "
-        "but got {} for element".format(int))
+        "Input type hint violation at TopMod: expected Tuple[TypeVariable[K], "
+        "TypeVariable[V]], got {}".format(int))
 
   def test_per_key_pipeline_checking_satisfied(self):
     d = (
@@ -2492,10 +2491,8 @@
 
     self.assertStartswith(
         e.exception.args[0],
-        "Type hint violation for 'CombinePerKey': "
-        "requires "
-        "Tuple[TypeVariable[K], Tuple[TypeVariable[K], TypeVariable[V]]] "
-        "but got Tuple[None, int] for element")
+        "Input type hint violation at ToDict: expected Tuple[TypeVariable[K], "
+        "TypeVariable[V]], got {}".format(int))
 
   def test_to_dict_pipeline_check_satisfied(self):
     d = (