How to keep a connection to external database in Cloud Dataflow

Question

I have an unbound dataflow pipeline that reads from Pub/Sub, applies a ParDo and writes to Cassandra. It applies only ParDo transformations so I am using the default global window with the default triggering even though the source is unbound.

In a pipeline like that how should I keep the connection to Cassandra?

Currently I am keeping it in startBundle like this:

private class CassandraWriter <T> extends DoFn<T, Void> {
  private transient Cluster cluster;
  private transient Session session;
  private transient MappingManager mappingManager;

  @Override
  public void startBundle(Context c) {
    this.cluster = Cluster.builder()
        .addContactPoints(hosts)
        .withPort(port)
        .withoutMetrics()
        .withoutJMXReporting()
        .build();
    this.session = cluster.connect(keyspace);
    this.mappingManager = new MappingManager(session);
  }

  @Override
  public void processElement(ProcessContext c) throws IOException {
    T element = c.element();
    Mapper<T> mapper = (Mapper<T>) mappingManager.mapper(element.getClass());
    mapper.save(element);
  }

  @Override
  public void finishBundle(Context c) throws IOException {
    session.close();
    cluster.close();
  }
}

However, this way a new connection is created for every element.

Another option is to pass it as a side input like in https://github.com/benjumanji/cassandra-dataflow:

public PDone apply(PCollection<T> input) {
  Pipeline p = input.getPipeline();

  CassandraWriteOperation<T> op = new CassandraWriteOperation<T>(this);

  Coder<CassandraWriteOperation<T>> coder =
    (Coder<CassandraWriteOperation<T>>)SerializableCoder.of(op.getClass());

  PCollection<CassandraWriteOperation<T>> opSingleton =
    p.apply(Create.<CassandraWriteOperation<T>>of(op)).setCoder(coder);

  final PCollectionView<CassandraWriteOperation<T>> opSingletonView =
    opSingleton.apply(View.<CassandraWriteOperation<T>>asSingleton());

  PCollection<Void> results = input.apply(ParDo.of(new DoFn<T, Void>() {
    @Override
    public void processElement(ProcessContext c) throws Exception {
       // use the side input here
    }
  }).withSideInputs(opSingletonView));

  PCollectionView<Iterable<Void>> voidView = results.apply(View.<Void>asIterable());

  opSingleton.apply(ParDo.of(new DoFn<CassandraWriteOperation<T>, Void>() {
    private static final long serialVersionUID = 0;

    @Override
    public void processElement(ProcessContext c) {
      CassandraWriteOperation<T> op = c.element();
      op.finalize();
    }

  }).withSideInputs(voidView));

  return new PDone();
}

However this way I have to use windowing since PCollectionView<Iterable<Void>> voidView = results.apply(View.<Void>asIterable()); applies a group by.

In general, how should a PTransform that writes from an unbounded PCollection to an external database keep its connection to the database?

Answer 1

You are correctly observing that a typical bundle size in the streaming/unbounded case is smaller compared to the batch/bounded case. The actual bundle size depends on many parameters, and sometimes bundles may contain a single element.

One way of solving this problem would be to use a pool of connections per worker, stored in a static state of your DoFn. You should be able to initialize it during the first call to startBundle, and use it across bundles. Alternatively, you can create a connection on demand and release it to the pool for reuse when no longer necessary.

You should make sure the static static is thread-safe, and that you aren't making any assumptions how Dataflow manages bundles.

Answer 2

As Davor Bonaci suggested, using a static variable solved the problem.

public class CassandraWriter<T> extends DoFn<T, Void> {
  private static final Logger log = LoggerFactory.getLogger(CassandraWriter.class);

  // Prevent multiple threads from creating multiple cluster connection in parallel.
  private static transient final Object lock = new Object();
  private static transient Cluster cluster;
  private static transient Session session;
  private static transient MappingManager mappingManager;

  private final String[] hosts;
  private final int port;
  private final String keyspace;

  public CassandraWriter(String[] hosts, int port, String keyspace) {
    this.hosts = hosts;
    this.port = port;
    this.keyspace = keyspace;
  }

  @Override
  public void startBundle(Context c) {
    synchronized (lock) {
      if (cluster == null) {
        cluster = Cluster.builder()
            .addContactPoints(hosts)
            .withPort(port)
            .withoutMetrics()
            .withoutJMXReporting()
            .build();
        session = cluster.connect(keyspace);
        mappingManager = new MappingManager(session);
      }
    }
  }

  @Override
  public void processElement(ProcessContext c) throws IOException {
    T element = c.element();
    Mapper<T> mapper = (Mapper<T>) mappingManager.mapper(element.getClass());
    mapper.save(element);
  }
}