All Downloads are FREE. Search and download functionalities are using the official Maven repository.

net.snowflake.ingest.streaming.internal.ChannelCache Maven / Gradle / Ivy

There is a newer version: 3.0.0
Show newest version
/*
 * Copyright (c) 2021-2024 Snowflake Computing Inc. All rights reserved.
 */

package net.snowflake.ingest.streaming.internal;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import net.snowflake.ingest.utils.ErrorCode;
import net.snowflake.ingest.utils.SFException;

/**
 * In-memory cache that stores the active channels for a given Streaming Ingest client, and the
 * channels belong to the same table will be stored together in order to combine the channel data
 * during flush. The key is a fully qualified table name and the value is a set of channels that
 * belongs to this table
 *
 * @param  type of column data ({@link ParquetChunkData})
 */
class ChannelCache {
  // Cache to hold all the valid channels, the key for the outer map is FullyQualifiedTableName and
  // the key for the inner map is ChannelName
  private final ConcurrentHashMap<
          String, ConcurrentHashMap>>
      cache = new ConcurrentHashMap<>();

  /** Flush information for each table including last flush time and if flush is needed */
  static class FlushInfo {
    final long lastFlushTime;
    final boolean needFlush;

    FlushInfo(long lastFlushTime, boolean needFlush) {
      this.lastFlushTime = lastFlushTime;
      this.needFlush = needFlush;
    }
  }

  /** Flush information for each table, only used when max chunks in blob is 1 */
  private final ConcurrentHashMap tableFlushInfo = new ConcurrentHashMap<>();

  /**
   * Add a channel to the channel cache
   *
   * @param channel
   */
  void addChannel(SnowflakeStreamingIngestChannelInternal channel) {
    ConcurrentHashMap> channels =
        this.cache.computeIfAbsent(
            channel.getFullyQualifiedTableName(), v -> new ConcurrentHashMap<>());

    // Update the last flush time for the table, add jitter to avoid all channels flush at the same
    // time when the blobs are not interleaved
    this.tableFlushInfo.putIfAbsent(
        channel.getFullyQualifiedTableName(), new FlushInfo(System.currentTimeMillis(), false));

    SnowflakeStreamingIngestChannelInternal oldChannel =
        channels.put(channel.getName(), channel);
    // Invalidate old channel if it exits to block new inserts and return error to users earlier
    if (oldChannel != null) {
      String invalidationCause =
          String.format("Old channel removed from cache, channelName=%s", channel.getName());
      oldChannel.invalidate("removed from cache", invalidationCause);
    }
  }

  /**
   * Get the last flush time for a table
   *
   * @param fullyQualifiedTableName fully qualified table name
   * @return last flush time in milliseconds
   */
  Long getLastFlushTime(String fullyQualifiedTableName) {
    FlushInfo tableFlushInfo = this.tableFlushInfo.get(fullyQualifiedTableName);
    if (tableFlushInfo == null) {
      throw new SFException(
          ErrorCode.INTERNAL_ERROR,
          String.format("Last flush time for table %s not found", fullyQualifiedTableName));
    }
    return tableFlushInfo.lastFlushTime;
  }

  /**
   * Set the last flush time for a table as the current time
   *
   * @param fullyQualifiedTableName fully qualified table name
   * @param lastFlushTime last flush time in milliseconds
   */
  void setLastFlushTime(String fullyQualifiedTableName, Long lastFlushTime) {
    this.tableFlushInfo.compute(
        fullyQualifiedTableName,
        (k, v) -> {
          if (v == null) {
            throw new SFException(
                ErrorCode.INTERNAL_ERROR,
                String.format("Last flush time for table %s not found", fullyQualifiedTableName));
          }
          return new FlushInfo(lastFlushTime, v.needFlush);
        });
  }

  /**
   * Get need flush flag for a table
   *
   * @param fullyQualifiedTableName fully qualified table name
   * @return need flush flag
   */
  boolean getNeedFlush(String fullyQualifiedTableName) {
    FlushInfo tableFlushInfo = this.tableFlushInfo.get(fullyQualifiedTableName);
    if (tableFlushInfo == null) {
      throw new SFException(
          ErrorCode.INTERNAL_ERROR,
          String.format("Need flush flag for table %s not found", fullyQualifiedTableName));
    }
    return tableFlushInfo.needFlush;
  }

  /**
   * Set need flush flag for a table
   *
   * @param fullyQualifiedTableName fully qualified table name
   * @param needFlush need flush flag
   */
  void setNeedFlush(String fullyQualifiedTableName, boolean needFlush) {
    this.tableFlushInfo.compute(
        fullyQualifiedTableName,
        (k, v) -> {
          if (v == null) {
            throw new SFException(
                ErrorCode.INTERNAL_ERROR,
                String.format("Need flush flag for table %s not found", fullyQualifiedTableName));
          }
          return new FlushInfo(v.lastFlushTime, needFlush);
        });
  }

  /** Returns an immutable set view of the mappings contained in the channel cache. */
  Set>>>
      entrySet() {
    return Collections.unmodifiableSet(cache.entrySet());
  }

  /** Returns an immutable set view of the keys contained in the channel cache. */
  Set keySet() {
    return Collections.unmodifiableSet(cache.keySet());
  }

  /** Close all channels in the channel cache */
  void closeAllChannels() {
    this.cache
        .values()
        .forEach(channels -> channels.values().forEach(channel -> channel.markClosed()));
  }

  /** Remove a channel in the channel cache if the channel sequencer matches */
  // TODO: background cleaner to cleanup old stale channels that are not closed?
  void removeChannelIfSequencersMatch(SnowflakeStreamingIngestChannelInternal channel) {
    cache.computeIfPresent(
        channel.getFullyQualifiedTableName(),
        (k, v) -> {
          SnowflakeStreamingIngestChannelInternal channelInCache = v.get(channel.getName());
          // We need to compare the channel sequencer in case the old channel was already been
          // removed
          return channelInCache != null
                  && channelInCache.getChannelSequencer() == channel.getChannelSequencer()
                  && v.remove(channel.getName()) != null
                  && v.isEmpty()
              ? null
              : v;
        });
  }

  /** Invalidate a channel in the channel cache if the channel sequencer matches */
  void invalidateChannelIfSequencersMatch(
      String dbName,
      String schemaName,
      String tableName,
      String channelName,
      Long channelSequencer,
      String invalidationCause) {
    String fullyQualifiedTableName = String.format("%s.%s.%s", dbName, schemaName, tableName);
    ConcurrentHashMap> channelsMapPerTable =
        cache.get(fullyQualifiedTableName);
    if (channelsMapPerTable != null) {
      SnowflakeStreamingIngestChannelInternal channel = channelsMapPerTable.get(channelName);
      if (channel != null && channel.getChannelSequencer().equals(channelSequencer)) {
        channel.invalidate("invalidate with matched sequencer", invalidationCause);
      }
    }
  }

  /** Get the number of key-value pairs in the cache */
  int getSize() {
    return cache.size();
  }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy