org.killbill.queue.OptimizePersistentBus Maven / Gradle / Ivy

Go to download
package org.killbill.queue;

import com.codahale.metrics.MetricRegistry;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.sql.DataSource;
import org.killbill.CreatorName;
import org.killbill.bus.DefaultPersistentBus;
import org.killbill.bus.api.BusEvent;
import org.killbill.bus.api.PersistentBusConfig;
import org.killbill.bus.dao.BusEventModelDao;
import org.killbill.bus.dao.PersistentBusSqlDao;
import org.killbill.clock.Clock;
import org.killbill.clock.DefaultClock;
import org.killbill.commons.jdbi.notification.DatabaseTransactionNotificationApi;
import org.killbill.queue.dao.QueueSqlDao;
import org.skife.config.ConfigurationObjectFactory;
import org.skife.jdbi.v2.DBI;
import org.skife.jdbi.v2.IDBI;
import org.skife.jdbi.v2.Transaction;
import org.slf4j.Logger;

/**
 *
 * @author joewee
 * @version 1.0.0
 * @date 2021/9/23 13:13
 */
public class OptimizePersistentBus extends DefaultPersistentBus  {

  private static final Logger log = org.slf4j.LoggerFactory.getLogger(OptimizePersistentBus.class);
  private final DBI dbi;
  private final AtomicBoolean isInitialized;

  public OptimizePersistentBus(IDBI dbi, Clock clock,
      PersistentBusConfig config,
      MetricRegistry metricRegistry,
      DatabaseTransactionNotificationApi databaseTransactionNotificationApi) {
    super(dbi, clock, config, metricRegistry, databaseTransactionNotificationApi);
    this.dbi = (DBI)dbi;
    this.isInitialized = new AtomicBoolean(false);
  }

  public OptimizePersistentBus(DataSource dataSource, Properties properties) {
    this(InTransaction.buildDDBI(dataSource), new DefaultClock(), new ConfigurationObjectFactory(properties).buildWithReplacements(PersistentBusConfig.class, ImmutableMap.of("instanceName", "main")),
        new MetricRegistry(), new DatabaseTransactionNotificationApi());
  }

  @Override
  public boolean initQueue() {
    if (isInitialized.compareAndSet(false, true)) {
      return super.initQueue();
    } else {
      return false;
    }
  }


  @Override
  public void stopQueue() {
    super.stopQueue();
    isInitialized.set(false);
  }


  public void post( Collection events ) throws EventBusException {
    try {
      if (isInitialized.get()) {
        final List> partition = Lists.partition(new ArrayList(events), 100);
        for (List list : partition) {
          List entries = new ArrayList<>(list.size());
          for (BusEvent event : list) {
            final String json = objectMapper.writeValueAsString(event);
            final BusEventModelDao entry = new BusEventModelDao(CreatorName.get(), super.getClock().getUTCNow(),
                event.getClass().getName(), json,
                event.getUserToken(), event.getSearchKey1(), event.getSearchKey2());
            entries.add(entry);
          }
          super.getDao().executeTransaction(
              (Transaction>) (transactional, status) -> {
                transactional.insertEntries(entries,
                    OptimizePersistentBus.super.getConfig().getTableName());
                return null;
              });
        }
      } else {
        log.warn("Attempting to post events in a non initialized bus");
      }
    } catch (final Exception e) {
      log.error("Failed to post BusEvent ", e);
    }
  }


  public void postFromTransaction(Collection events, final Connection connection) throws EventBusException {
    if (!isInitialized.get()) {
      log.warn("Attempting to post events in a non initialized bus");
      return;
    }
    //将events按100的大小分批次提交
    final List> partition = Lists.partition(new ArrayList(events), 100);
    for (List list : partition) {
      List entries = new ArrayList<>(list.size());
      for (BusEvent event : list) {
        final String json;
        try {
          json = objectMapper.writeValueAsString(event);
        } catch (final JsonProcessingException e) {
          log.error("Unable to serialize event " + event, e);
          return;
        }
        final BusEventModelDao entry = new BusEventModelDao(CreatorName.get(),
            super.getClock().getUTCNow(),
            event.getClass().getName(),
            json,
            event.getUserToken(),
            event.getSearchKey1(),
            event.getSearchKey2());
        entries.add(entry);
      }
      final InTransaction.InTransactionHandler handler = transactional -> {
        transactional.insertEntries(entries,
            OptimizePersistentBus.super.getConfig().getTableName());
        return null;
      };
      InTransaction.execute(dbi, connection, handler, PersistentBusSqlDao.class);
    }
  }
}