2021-05-26

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.statefun.flink.core.StatefulFunctionsConfig;
import org.apache.flink.statefun.flink.core.message.MessageFactoryType;
import org.apache.flink.statefun.flink.core.message.RoutableMessage;
import org.apache.flink.statefun.flink.core.message.RoutableMessageBuilder;
import org.apache.flink.statefun.flink.datastream.StatefulFunctionDataStreamBuilder;
import org.apache.flink.statefun.flink.datastream.StatefulFunctionEgressStreams;
import org.apache.flink.statefun.sdk.Context;
import org.apache.flink.statefun.sdk.FunctionType;
import org.apache.flink.statefun.sdk.StatefulFunction;
import org.apache.flink.statefun.sdk.annotations.Persisted;
import org.apache.flink.statefun.sdk.io.EgressIdentifier;
import org.apache.flink.statefun.sdk.state.PersistedValue;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.PrintSinkFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import javax.annotation.Nullable;
import java.net.URI;
import java.time.Duration;
import java.util.concurrent.ThreadLocalRandom;

import static org.apache.flink.statefun.flink.datastream.RequestReplyFunctionBuilder.requestReplyFunctionBuilder;

/**
 * @author chenzhuoyu
 * @date 2021/5/26 17:48
 */
public class test {
    private static EgressIdentifier<String> GREETINGS = new EgressIdentifier<>("example", "greetings", String.class);


    public static void main(String[] args) throws Exception {


        FunctionType GREET = new FunctionType("example", "greet");
        FunctionType REMOTE_GREET = new FunctionType("example", "remote-greet");
//        FunctionType REMOTE_GREET = new FunctionType("example", "remote-greet");

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        StatefulFunctionsConfig statefunConfig = StatefulFunctionsConfig.fromEnvironment(env);
        statefunConfig.setGlobalConfiguration("someGlobalConfig", "foobar");
        statefunConfig.setFactoryType(MessageFactoryType.WITH_KRYO_PAYLOADS);

        DataStream<String> names = env.addSource(new NameSource());

        DataStream<RoutableMessage> namesIngress =
                names.map(name ->
                        RoutableMessageBuilder.builder()
                                .withTargetAddress(new FunctionType("example", "greet"), name)
                                .withMessageBody(name)
                                .build()
                );

        StatefulFunctionEgressStreams egresses =
                StatefulFunctionDataStreamBuilder.builder("example")
                        .withDataStreamAsIngress(namesIngress)
                        .withFunctionProvider(GREET,unused -> new MyFunction())
//                        .withRequestReplyRemoteFunction(
//                                requestReplyFunctionBuilder(
//                                        REMOTE_GREET, URI.create("http://localhost:5000/statefun"))
//                                        .withPersistedState("seen_count")
//                                        .withMaxRequestDuration(Duration.ofSeconds(15))
//                                        .withMaxNumBatchRequests(500))
                        .withEgressId(GREETINGS)
                        .build(env);



        DataStream<String> output = egresses.getDataStreamForEgressId(GREETINGS);
//
//    // -----------------------------------------------------------------------------------------
//    // the rest of the pipeline
//    // -----------------------------------------------------------------------------------------
//
        output.map(
                new RichMapFunction<String, String>() {
                    @Override
                    public String map(String value) {
                        return "'" + value + "'";
                    }
                })
                .addSink(new PrintSinkFunction<>());

        env.execute();


    }

    private static class MyFunction implements StatefulFunction {
        @Persisted
        private final PersistedValue<Integer> seenCount = PersistedValue.of("seen", Integer.class);
        @Override
        public void invoke(Context context, Object input) {
            int seen = seenCount.updateAndGet(MyFunction::increment);
            context.send(GREETINGS, String.format("Hello %s at the %d-th time", input, seen));

        }
        private static int increment(@Nullable Integer n) {
            return n == null ? 1 : n + 1;
        }
    }
    private static final class NameSource implements SourceFunction<String> {

        private static final long serialVersionUID = 1;

        private volatile boolean canceled;

        @Override
        public void run(SourceContext<String> ctx) throws InterruptedException {
            String[] names = {"Stephan", "Igal", "Gordon", "Seth", "Marta"};
            ThreadLocalRandom random = ThreadLocalRandom.current();
            while (true) {
                int index = random.nextInt(names.length);
                final String name = names[index];
                synchronized (ctx.getCheckpointLock()) {
                    if (canceled) {
                        return;
                    }
                    ctx.collect(name);
                }
                Thread.sleep(1000);
            }
        }

        @Override
        public void cancel() {
            canceled = true;
        }
    }
}