Fix pulsar sink and source state (#5046)

(cherry picked from commit 40d6248)

pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/ContextImpl.java

@@ -18,12 +18,15 @@
  */
 package org.apache.pulsar.functions.instance;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.gson.Gson;
 import com.google.gson.reflect.TypeToken;
+import io.netty.buffer.ByteBuf;
 import io.prometheus.client.CollectorRegistry;
 import io.prometheus.client.Summary;
 import lombok.Getter;
 import lombok.Setter;
+import org.apache.bookkeeper.api.kv.Table;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.pulsar.client.api.*;
 import org.apache.pulsar.client.impl.ProducerBuilderImpl;
@@ -47,7 +50,6 @@
 import java.nio.ByteBuffer;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 
 import static com.google.common.base.Preconditions.checkState;
@@ -72,9 +74,8 @@ class ContextImpl implements Context, SinkContext, SourceContext {
     private final SecretsProvider secretsProvider;
     private final Map<String, Object> secretsMap;
 
-    @Getter
-    @Setter
-    private StateContextImpl stateContext;
+    @VisibleForTesting
+    StateContextImpl stateContext;
     private Map<String, Object> userConfigs;
 
     private ComponentStatsManager statsManager;
@@ -95,7 +96,8 @@ class ContextImpl implements Context, SinkContext, SourceContext {
 
     public ContextImpl(InstanceConfig config, Logger logger, PulsarClient client,
                        SecretsProvider secretsProvider, CollectorRegistry collectorRegistry, String[] metricsLabels,
-                       Function.FunctionDetails.ComponentType componentType, ComponentStatsManager statsManager) {
+                       Function.FunctionDetails.ComponentType componentType, ComponentStatsManager statsManager,
+                       Table<ByteBuf, ByteBuf> stateTable) {
         this.config = config;
         this.logger = logger;
         this.publishProducers = new HashMap<>();
@@ -146,6 +148,10 @@ public ContextImpl(InstanceConfig config, Logger logger, PulsarClient client,
                 .quantile(0.999, 0.01)
                 .register(collectorRegistry);
         this.componentType = componentType;
+
+        if (null != stateTable) {
+            this.stateContext = new StateContextImpl(stateTable);
+        }
     }
 
     public void setCurrentMessageContext(Record<?> record) {

pulsar-functions/instance/src/main/java/org/apache/pulsar/functions/instance/JavaInstanceRunnable.java

@@ -213,7 +213,7 @@ ContextImpl setupContext() {
         Logger instanceLog = LoggerFactory.getLogger(
                 "function-" + instanceConfig.getFunctionDetails().getName());
         return new ContextImpl(instanceConfig, instanceLog, client, secretsProvider,
-                collectorRegistry, metricsLabels, this.componentType, this.stats);
+                collectorRegistry, metricsLabels, this.componentType, this.stats, stateTable);
     }
 
     /**
@@ -232,10 +232,6 @@ public void run() {
                     this.componentType);
 
             javaInstance = setupJavaInstance();
-            if (null != stateTable) {
-                StateContextImpl stateContext = new StateContextImpl(stateTable);
-                javaInstance.getContext().setStateContext(stateContext);
-            }
             while (true) {
                 currentRecord = readInput();
 

pulsar-functions/instance/src/test/java/org/apache/pulsar/functions/instance/ContextImplTest.java

@@ -36,6 +36,7 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
+import org.apache.bookkeeper.api.kv.Table;
 import org.apache.pulsar.client.api.Producer;
 import org.apache.pulsar.client.api.Schema;
 import org.apache.pulsar.client.api.TypedMessageBuilder;
@@ -81,13 +82,12 @@ public void setup() {
         TypedMessageBuilder messageBuilder = spy(new TypedMessageBuilderImpl(mock(ProducerBase.class), Schema.STRING));
         doReturn(new CompletableFuture<>()).when(messageBuilder).sendAsync();
         when(producer.newMessage()).thenReturn(messageBuilder);
-
         context = new ContextImpl(
             config,
             logger,
             client,
             new EnvironmentBasedSecretsProvider(), new CollectorRegistry(), new String[0],
-                FunctionDetails.ComponentType.FUNCTION, null);
+                FunctionDetails.ComponentType.FUNCTION, null, null);
         context.setCurrentMessageContext((Record<String>) () -> null);
     }
 
@@ -98,6 +98,7 @@ public void testIncrCounterStateDisabled() {
 
     @Test(expectedExceptions = IllegalStateException.class)
     public void testGetCounterStateDisabled() {
+
         context.getCounter("test-key");
     }
 
@@ -113,35 +114,31 @@ public void testGetStateStateDisabled() {
 
     @Test
     public void testIncrCounterStateEnabled() throws Exception {
-        StateContextImpl stateContext = mock(StateContextImpl.class);
-        context.setStateContext(stateContext);
+        context.stateContext = mock(StateContextImpl.class);
         context.incrCounterAsync("test-key", 10L);
-        verify(stateContext, times(1)).incrCounter(eq("test-key"), eq(10L));
+        verify(context.stateContext, times(1)).incrCounter(eq("test-key"), eq(10L));
     }
 
     @Test
     public void testGetCounterStateEnabled() throws Exception {
-        StateContextImpl stateContext = mock(StateContextImpl.class);
-        context.setStateContext(stateContext);
+        context.stateContext = mock(StateContextImpl.class);
         context.getCounterAsync("test-key");
-        verify(stateContext, times(1)).getCounter(eq("test-key"));
+        verify(context.stateContext, times(1)).getCounter(eq("test-key"));
     }
 
     @Test
     public void testPutStateStateEnabled() throws Exception {
-        StateContextImpl stateContext = mock(StateContextImpl.class);
-        context.setStateContext(stateContext);
+        context.stateContext = mock(StateContextImpl.class);
         ByteBuffer buffer = ByteBuffer.wrap("test-value".getBytes(UTF_8));
         context.putStateAsync("test-key", buffer);
-        verify(stateContext, times(1)).put(eq("test-key"), same(buffer));
+        verify(context.stateContext, times(1)).put(eq("test-key"), same(buffer));
     }
 
     @Test
     public void testGetStateStateEnabled() throws Exception {
-        StateContextImpl stateContext = mock(StateContextImpl.class);
-        context.setStateContext(stateContext);
+        context.stateContext = mock(StateContextImpl.class);
         context.getStateAsync("test-key");
-        verify(stateContext, times(1)).get(eq("test-key"));
+        verify(context.stateContext, times(1)).get(eq("test-key"));
     }
 
     @Test

tests/docker-images/java-test-functions/pom.xml

@@ -29,6 +29,13 @@
     <groupId>org.apache.pulsar.tests</groupId>
     <artifactId>java-test-functions</artifactId>
     <name>Apache Pulsar :: Tests :: Docker Images :: Java Test Functions</name>
+    <dependencies>
+        <dependency>
+            <groupId>org.apache.pulsar</groupId>
+            <artifactId>pulsar-io-core</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+    </dependencies>
     <packaging>jar</packaging>
 
     <profiles>

tests/docker-images/java-test-functions/src/main/java/org/apache/pulsar/tests/integration/io/TestStateSink.java

@@ -0,0 +1,51 @@
+/**
+ * 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.pulsar.tests.integration.io;
+
+import org.apache.pulsar.functions.api.Record;
+import org.apache.pulsar.io.core.Sink;
+import org.apache.pulsar.io.core.SinkContext;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+public class TestStateSink implements Sink<String> {
+
+    private SinkContext sinkContext;
+    private int count;
+
+    @Override
+    public void open(Map<String, Object> config, SinkContext sinkContext) throws Exception {
+        sinkContext.putState("initial", ByteBuffer.wrap("val1".getBytes()));
+        this.sinkContext = sinkContext;
+    }
+
+    @Override
+    public void write(Record<String> record) throws Exception {
+        String initial = new String(sinkContext.getState("initial").array());
+        String val = String.format("%s-%d", initial, count);
+        sinkContext.putState("now", ByteBuffer.wrap(val.getBytes()));
+        count++;
+    }
+
+    @Override
+    public void close() throws Exception {
+
+    }
+}

tests/docker-images/java-test-functions/src/main/java/org/apache/pulsar/tests/integration/io/TestStateSource.java

@@ -0,0 +1,54 @@
+/**
+ * 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.pulsar.tests.integration.io;
+
+import org.apache.pulsar.functions.api.Record;
+import org.apache.pulsar.io.core.Source;
+import org.apache.pulsar.io.core.SourceContext;
+
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+public class TestStateSource implements Source<String> {
+
+
+    private SourceContext sourceContext;
+    private int count;
+
+    @Override
+    public void open(Map<String, Object> config, SourceContext sourceContext) throws Exception {
+        sourceContext.putState("initial", ByteBuffer.wrap("val1".getBytes()));
+        this.sourceContext = sourceContext;
+    }
+
+    @Override
+    public Record<String> read() throws Exception {
+        Thread.sleep(50);
+        String initial = new String(sourceContext.getState("initial").array());
+        String val = String.format("%s-%d", initial, count);
+        sourceContext.putState("now", ByteBuffer.wrap(val.getBytes()));
+        count++;
+        return () -> val;
+    }
+
+    @Override
+    public void close() throws Exception {
+
+    }
+}