samza-core/src/main/java/org/apache/samza/operators/impl/WindowOperatorImpl.java - samza - Git at Google

 // CHECKSTYLE:OFF
 // Turn off checkstyle for this class because of a checkstyle bug in handling nested typed collections
 /*
  * 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.samza.operators.impl;

 import com.google.common.base.Preconditions;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.ConcurrentHashMap;
 import org.apache.samza.context.Context;
 import org.apache.samza.operators.functions.FoldLeftFunction;
 import org.apache.samza.operators.functions.MapFunction;
 import org.apache.samza.operators.functions.SupplierFunction;
 import org.apache.samza.operators.impl.store.TimeSeriesKey;
 import org.apache.samza.operators.impl.store.TimeSeriesStore;
 import org.apache.samza.operators.impl.store.TimeSeriesStoreImpl;
 import org.apache.samza.operators.spec.OperatorSpec;
 import org.apache.samza.operators.spec.WindowOperatorSpec;
 import org.apache.samza.operators.triggers.FiringType;
 import org.apache.samza.operators.triggers.RepeatingTriggerImpl;
 import org.apache.samza.operators.triggers.TimeTrigger;
 import org.apache.samza.operators.triggers.Trigger;
 import org.apache.samza.operators.triggers.TriggerImpl;
 import org.apache.samza.operators.triggers.TriggerImpls;
 import org.apache.samza.operators.windows.AccumulationMode;
 import org.apache.samza.operators.windows.WindowKey;
 import org.apache.samza.operators.windows.WindowPane;
 import org.apache.samza.operators.windows.internal.WindowInternal;
 import org.apache.samza.operators.windows.internal.WindowType;
 import org.apache.samza.storage.kv.ClosableIterator;
 import org.apache.samza.storage.kv.KeyValueStore;
 import org.apache.samza.task.MessageCollector;
 import org.apache.samza.task.TaskCoordinator;
 import org.apache.samza.util.Clock;
 import org.apache.samza.util.TimestampedValue;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;

 /**
  * Implementation of a window operator that groups messages into finite windows for processing.
  *
  * This class implements the processing logic for various types of windows and triggers. It tracks and manages state for
  * all open windows, the active triggers that correspond to each of the windows and the pending callbacks. It provides
  * an implementation of {@link TriggerScheduler} that {@link TriggerImpl}s can use
  * to schedule and cancel callbacks. It also orchestrates the flow of messages through the various
  * {@link TriggerImpl}s.
  *
  * <p> An instance of a {@link TriggerImplHandler} is created corresponding to each {@link Trigger}
  * configured for a particular window. For every message added to the window, this class looks up the corresponding
  * {@link TriggerImplHandler} for the trigger and invokes {@link TriggerImplHandler#onMessage(TriggerKey, Object,
  * MessageCollector, TaskCoordinator)} The {@link TriggerImplHandler} maintains the {@link TriggerImpl} instance along
  * with whether it has been canceled yet or not. Then, the {@link TriggerImplHandler} invokes onMessage on underlying
  * its {@link TriggerImpl} instance. A {@link TriggerImpl} instance is scoped to a window and its firing determines when
  * results for its window are emitted.
  *
  * The {@link WindowOperatorImpl} checks if the trigger fired and returns the result of the firing.
  *
  * @param <M> the type of the incoming message
  * @param <K> the type of the key in the incoming message
  *
  */
 public class WindowOperatorImpl<M, K> extends OperatorImpl<M, WindowPane<K, Object>> {
   // Object == Collection<M> || WV
   private static final Logger LOG = LoggerFactory.getLogger(WindowOperatorImpl.class);

   private final WindowOperatorSpec<M, K, Object> windowOpSpec;
   private final Clock clock;
   private final WindowInternal<M, K, Object> window;
   private final FoldLeftFunction<M, Object> foldLeftFn;
   private final SupplierFunction<Object> initializer;
   private final MapFunction<M, K> keyFn;

   private final TriggerScheduler<K> triggerScheduler;
   private final Map<TriggerKey<K>, TriggerImplHandler> triggers = new ConcurrentHashMap<>();
   private TimeSeriesStore<K, Object> timeSeriesStore;

   public WindowOperatorImpl(WindowOperatorSpec<M, K, Object> windowOpSpec, Clock clock) {
     this.windowOpSpec = windowOpSpec;
     this.clock = clock;
     this.window = windowOpSpec.getWindow();
     this.foldLeftFn = window.getFoldLeftFunction();
     this.initializer = window.getInitializer();
     this.keyFn = window.getKeyExtractor();
     this.triggerScheduler= new TriggerScheduler(clock);
   }

   @Override
   protected void handleInit(Context context) {

     KeyValueStore<TimeSeriesKey<K>, Object> store =
         (KeyValueStore<TimeSeriesKey<K>, Object>) context.getTaskContext().getStore(windowOpSpec.getOpId());

     if (initializer != null) {
       initializer.init(context);
     }

     if (keyFn != null) {
       keyFn.init(context);
     }

     // For aggregating windows, we use the store in over-write mode since we only retain the aggregated
     // value. Else, we use the store in append-mode.
     if (foldLeftFn != null) {
       foldLeftFn.init(context);
       timeSeriesStore = new TimeSeriesStoreImpl(store, false);
     } else {
       timeSeriesStore = new TimeSeriesStoreImpl(store, true);
     }
   }

   @Override
   protected CompletionStage<Collection<WindowPane<K, Object>>> handleMessageAsync(M message, MessageCollector collector,
       TaskCoordinator coordinator) {
     LOG.trace("Processing message envelope: {}", message);
     List<WindowPane<K, Object>> results = new ArrayList<>();

     K key = (keyFn != null) ? keyFn.apply(message) : null;
     long timestamp = getWindowTimestamp(message);

     // For aggregating windows, we only store the aggregated window value.
     // For non-aggregating windows, we store all messages in the window.
     if (foldLeftFn == null) {
       timeSeriesStore.put(key, message, timestamp); // store is in append mode
     } else {
       List<Object> existingState = getValues(key, timestamp);

       Preconditions.checkState(existingState.size() <= 1, String.format("WindowState for aggregating windows " +
           "must not contain more than one entry per window. Current size: %s", existingState.size()));

       if (existingState.size() == 0) {
         LOG.trace("No existing state found for key {} Invoking initializer.", key);
       }
       Object oldVal = existingState.size() == 0 ? initializer.get() : existingState.get(0);
       Object aggregatedValue = foldLeftFn.apply(message, oldVal);

       timeSeriesStore.put(key, aggregatedValue, timestamp); // store is in over-write mode
     }

     if (window.getEarlyTrigger() != null) {
       TriggerKey<K> triggerKey = new TriggerKey<>(FiringType.EARLY, key, timestamp);
       TriggerImplHandler triggerImplHandler = getOrCreateTriggerImplHandler(triggerKey, window.getEarlyTrigger());
       Optional<WindowPane<K, Object>> maybeTriggeredPane =
           triggerImplHandler.onMessage(triggerKey, message, collector, coordinator);
       maybeTriggeredPane.ifPresent(results::add);
     }

     if (window.getDefaultTrigger() != null) {
       TriggerKey<K> triggerKey = new TriggerKey<>(FiringType.DEFAULT, key, timestamp);
       TriggerImplHandler triggerImplHandler = getOrCreateTriggerImplHandler(triggerKey, window.getDefaultTrigger());
       Optional<WindowPane<K, Object>> maybeTriggeredPane =
           triggerImplHandler.onMessage(triggerKey, message, collector, coordinator);
       maybeTriggeredPane.ifPresent(results::add);
     }

     return CompletableFuture.completedFuture(results);
   }

   @Override
   public Collection<WindowPane<K, Object>> handleTimer(MessageCollector collector, TaskCoordinator coordinator) {
     LOG.trace("Processing time triggers");
     List<WindowPane<K, Object>> results = new ArrayList<>();
     List<TriggerKey<K>> keys = triggerScheduler.runPendingCallbacks();

     for (TriggerKey<K> key : keys) {
       TriggerImplHandler triggerImplHandler = triggers.get(key);
       if (triggerImplHandler != null) {
         Optional<WindowPane<K, Object>> maybeTriggeredPane = triggerImplHandler.onTimer(key, collector, coordinator);
         maybeTriggeredPane.ifPresent(results::add);
       }
     }
     LOG.trace("Triggered panes: " + results.size());
     return results;
   }

   @Override
   protected OperatorSpec<M, WindowPane<K, Object>> getOperatorSpec() {
     return windowOpSpec;
   }

   @Override
   protected Collection<WindowPane<K, Object>> handleEndOfStream(MessageCollector collector, TaskCoordinator coordinator) {
     List<WindowPane<K, Object>> results = new ArrayList<>();
     Set<TriggerKey<K>> triggerKeys = new HashSet<>(triggers.keySet());
     for(TriggerKey<K> triggerKey : triggerKeys) {
       Optional<WindowPane<K, Object>> triggerResult = onTriggerFired(triggerKey, collector, coordinator);
       triggerResult.ifPresent(results::add);
     }
     return results;
   }

   @Override
   protected void handleClose() {
     if (foldLeftFn != null) {
       foldLeftFn.close();
     }
     if (timeSeriesStore != null) {
       timeSeriesStore.close();
     }
     if (initializer != null) {
       initializer.close();
     }
     if (keyFn != null) {
       keyFn.close();
     }
   }

   private TriggerImplHandler getOrCreateTriggerImplHandler(TriggerKey<K> triggerKey, Trigger<M> trigger) {
     TriggerImplHandler wrapper = triggers.get(triggerKey);
     if (wrapper != null) {
       LOG.trace("Returning existing trigger wrapper for {}", triggerKey);
       return wrapper;
     }

     LOG.trace("Creating a new trigger wrapper for {}", triggerKey);

     TriggerImpl<M, K> triggerImpl = TriggerImpls.createTriggerImpl(trigger, clock, triggerKey);
     wrapper = new TriggerImplHandler(triggerKey, triggerImpl);
     triggers.put(triggerKey, wrapper);

     return wrapper;
   }

   /**
    * Handles trigger firings and returns the optional result.
    */
   private Optional<WindowPane<K, Object>> onTriggerFired(TriggerKey<K> triggerKey, MessageCollector collector,
       TaskCoordinator coordinator) {
     LOG.trace("Trigger key {} fired." , triggerKey);

     TriggerImplHandler wrapper = triggers.get(triggerKey);
     long timestamp = triggerKey.getTimestamp();
     K key = triggerKey.getKey();
     List<Object> existingState = getValues(key, timestamp);

     if (existingState == null || existingState.size() == 0) {
       LOG.trace("No state found for triggerKey: {}", triggerKey);
       return Optional.empty();
     }

     Object windowVal = window.getFoldLeftFunction() == null ? existingState : existingState.get(0);

     WindowPane<K, Object> paneOutput = computePaneOutput(triggerKey, windowVal);

     // Handle different accumulation modes.
     if (window.getAccumulationMode() == AccumulationMode.DISCARDING) {
       LOG.trace("Clearing state for trigger key: {}", triggerKey);
       timeSeriesStore.remove(key, timestamp);
     }

     // Cancel all early triggers too when the default trigger fires. Also, clean all state for the key.
     // note: We don't handle late arrivals yet, So, all arrivals are either early or on-time.
     if (triggerKey.getType() == FiringType.DEFAULT) {

       LOG.trace("Default trigger fired. Canceling triggers for {}", triggerKey);

       cancelTrigger(triggerKey, true);
       cancelTrigger(new TriggerKey(FiringType.EARLY, triggerKey.getKey(), triggerKey.getTimestamp()), true);
       timeSeriesStore.remove(key, timestamp);
     }

     // Cancel non-repeating early triggers. All early triggers should be removed from the "triggers" map only after the
     // firing of their corresponding default trigger. Removing them pre-maturely (immediately after cancellation) will
     // will create a new {@link TriggerImplWrapper} instance at a future invocation of getOrCreateTriggerWrapper().
     // This would cause an already canceled trigger to fire again for the window.

     if (triggerKey.getType() == FiringType.EARLY && !wrapper.isRepeating()) {
       cancelTrigger(triggerKey, false);
     }

     return Optional.of(paneOutput);
   }

   /**
    * Computes the pane output corresponding to a {@link TriggerKey} that fired.
    */
   private WindowPane<K, Object> computePaneOutput(TriggerKey<K> triggerKey, Object windowVal) {
     WindowKey<K> windowKey = new WindowKey(triggerKey.getKey(), Long.toString(triggerKey.getTimestamp()));
     WindowPane<K, Object> paneOutput =
         new WindowPane<>(windowKey, windowVal, window.getAccumulationMode(), triggerKey.getType());
     LOG.trace("Emitting pane output for trigger key {}", triggerKey);
     return paneOutput;
   }

   /**
    * Cancels the firing of the {@link TriggerImpl} identified by this {@link TriggerKey} and optionally removes it.
    */
   private void cancelTrigger(TriggerKey<K> triggerKey, boolean shouldRemove) {
     TriggerImplHandler triggerImplHandler = triggers.get(triggerKey);
     if (triggerImplHandler != null) {
       triggerImplHandler.cancel();
     }
     if (shouldRemove && triggerKey != null) {
       triggers.remove(triggerKey);
     }
   }

   /**
    * Computes the timestamp of the window this message should belong to.
    *
    * In the case of tumbling windows, timestamp of a window is defined as the start timestamp of its corresponding window
    * interval. For instance, if the tumbling interval is 10 seconds, all messages that arrive between [1000, 1010]
    * are assigned to the window with timestamp "1000"
    *
    * In the case of session windows, timestamp is defined as the timestamp of the earliest message in the window.
    * For instance, if the session gap is 10 seconds, and the first message in the window arrives at "1002" seconds,
    * all messages (that arrive within 10 seconds of their previous message) are assigned a timestamp "1002".
    *
    * @param message the input message
    * @return the timestamp of the window this message should belong to
    */
   private long getWindowTimestamp(M message) {
     if (window.getWindowType() == WindowType.TUMBLING) {
       long triggerDurationMs = ((TimeTrigger<M>) window.getDefaultTrigger()).getDuration().toMillis();
       final long now = clock.currentTimeMillis();
       // assign timestamp to be the start timestamp of the window boundary
       long timestamp = now - now % triggerDurationMs;
       return timestamp;
     } else {
       K key = keyFn.apply(message);
       // get the value with the earliest timestamp for the provided key.
       ClosableIterator<TimestampedValue<Object>> iterator = timeSeriesStore.get(key, 0, Long.MAX_VALUE, 1);
       List<TimestampedValue<Object>> timestampedValues = toList(iterator);

       // If there are no existing sessions for the key, we return the current timestamp. If not, return the
       // timestamp of the earliest message.
       long timestamp = (timestampedValues.isEmpty())? clock.currentTimeMillis() : timestampedValues.get(0).getTimestamp();

       return timestamp;
     }
   }

   /**
    * Return a list of values in the store for the provided key and timestamp
    *
    * @param key the key to look up in the store
    * @param timestamp the timestamp to look up in the store
    * @return the list of values for the provided key
    */
   private List<Object> getValues(K key, long timestamp) {
     ClosableIterator<TimestampedValue<Object>> iterator = timeSeriesStore.get(key, timestamp);
     List<TimestampedValue<Object>> timestampedValues = toList(iterator);
     List<Object> values = timestampedValues.stream().map(element -> element.getValue()).collect(Collectors.toList());

     LOG.trace("Returning {} for key {} and timestamp {}", new Object[] {values, key, timestamp});
     return values;
   }

   /**
    * Returns an unmodifiable list of all elements in the provided iterator.
    * The iterator is guaranteed to be closed after its execution.
    *
    * @param iterator the provided iterator.
    * @param <V> the type of elements in the iterator
    * @return a list of all elements returned by the iterator
    */
   static <V>  List<V> toList(ClosableIterator<V> iterator) {
     Preconditions.checkNotNull(iterator);

     List<V> values = new ArrayList<>();
     try {
       while (iterator.hasNext()) {
         values.add(iterator.next());
       }
     } finally {
         iterator.close();
     }
     return Collections.unmodifiableList(values);
   }

   /**
    * State corresponding to a created {@link TriggerImpl} instance.
    */
   private class TriggerImplHandler {
     // The context, and the {@link TriggerImpl} instance corresponding to this triggerKey
     private final TriggerImpl<M, K> impl;
     // Guard to ensure that we don't invoke onMessage or onTimer on already cancelled triggers
     private boolean isCancelled = false;

     public TriggerImplHandler(TriggerKey<K> key, TriggerImpl<M, K> impl) {
       this.impl = impl;
     }

     public Optional<WindowPane<K, Object>> onMessage(TriggerKey<K> triggerKey, M message,
         MessageCollector collector, TaskCoordinator coordinator) {
       if (!isCancelled) {
         LOG.trace("Forwarding callbacks for {}", message);
         impl.onMessage(message, triggerScheduler);

         if (impl.shouldFire()) {
           // repeating trigger can trigger multiple times, So, clear the state to allow future triggerings.
           if (impl instanceof RepeatingTriggerImpl) {
             ((RepeatingTriggerImpl<M, K>) impl).clear();
           }
           return onTriggerFired(triggerKey, collector, coordinator);
         }
       }
       return Optional.empty();
     }

     public Optional<WindowPane<K, Object>> onTimer(TriggerKey<K> key, MessageCollector collector,
         TaskCoordinator coordinator) {
       if (impl.shouldFire() && !isCancelled) {
         LOG.trace("Triggering timer triggers");

         // repeating trigger can trigger multiple times, So, clear the trigger to allow future triggerings.
         if (impl instanceof RepeatingTriggerImpl) {
           ((RepeatingTriggerImpl<M, K>) impl).clear();
         }
         return onTriggerFired(key, collector, coordinator);
       }
       return Optional.empty();
     }

     public void cancel() {
       impl.cancel();
       isCancelled = true;
     }

     public boolean isRepeating() {
       return this.impl instanceof RepeatingTriggerImpl;
     }
   }
 }
	// CHECKSTYLE:OFF
	// Turn off checkstyle for this class because of a checkstyle bug in handling nested typed collections
	/*
	* 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.samza.operators.impl;

	import com.google.common.base.Preconditions;
	import java.util.concurrent.CompletableFuture;
	import java.util.concurrent.CompletionStage;
	import java.util.concurrent.ConcurrentHashMap;
	import org.apache.samza.context.Context;
	import org.apache.samza.operators.functions.FoldLeftFunction;
	import org.apache.samza.operators.functions.MapFunction;
	import org.apache.samza.operators.functions.SupplierFunction;
	import org.apache.samza.operators.impl.store.TimeSeriesKey;
	import org.apache.samza.operators.impl.store.TimeSeriesStore;
	import org.apache.samza.operators.impl.store.TimeSeriesStoreImpl;
	import org.apache.samza.operators.spec.OperatorSpec;
	import org.apache.samza.operators.spec.WindowOperatorSpec;
	import org.apache.samza.operators.triggers.FiringType;
	import org.apache.samza.operators.triggers.RepeatingTriggerImpl;
	import org.apache.samza.operators.triggers.TimeTrigger;
	import org.apache.samza.operators.triggers.Trigger;
	import org.apache.samza.operators.triggers.TriggerImpl;
	import org.apache.samza.operators.triggers.TriggerImpls;
	import org.apache.samza.operators.windows.AccumulationMode;
	import org.apache.samza.operators.windows.WindowKey;
	import org.apache.samza.operators.windows.WindowPane;
	import org.apache.samza.operators.windows.internal.WindowInternal;
	import org.apache.samza.operators.windows.internal.WindowType;
	import org.apache.samza.storage.kv.ClosableIterator;
	import org.apache.samza.storage.kv.KeyValueStore;
	import org.apache.samza.task.MessageCollector;
	import org.apache.samza.task.TaskCoordinator;
	import org.apache.samza.util.Clock;
	import org.apache.samza.util.TimestampedValue;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.Collections;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Optional;
	import java.util.Set;
	import java.util.stream.Collectors;

	/**
	* Implementation of a window operator that groups messages into finite windows for processing.
	*
	* This class implements the processing logic for various types of windows and triggers. It tracks and manages state for
	* all open windows, the active triggers that correspond to each of the windows and the pending callbacks. It provides
	* an implementation of {@link TriggerScheduler} that {@link TriggerImpl}s can use
	* to schedule and cancel callbacks. It also orchestrates the flow of messages through the various
	* {@link TriggerImpl}s.
	*
	* <p> An instance of a {@link TriggerImplHandler} is created corresponding to each {@link Trigger}
	* configured for a particular window. For every message added to the window, this class looks up the corresponding
	* {@link TriggerImplHandler} for the trigger and invokes {@link TriggerImplHandler#onMessage(TriggerKey, Object,
	* MessageCollector, TaskCoordinator)} The {@link TriggerImplHandler} maintains the {@link TriggerImpl} instance along
	* with whether it has been canceled yet or not. Then, the {@link TriggerImplHandler} invokes onMessage on underlying
	* its {@link TriggerImpl} instance. A {@link TriggerImpl} instance is scoped to a window and its firing determines when
	* results for its window are emitted.
	*
	* The {@link WindowOperatorImpl} checks if the trigger fired and returns the result of the firing.
	*
	* @param <M> the type of the incoming message
	* @param <K> the type of the key in the incoming message
	*
	*/
	public class WindowOperatorImpl<M, K> extends OperatorImpl<M, WindowPane<K, Object>> {
	// Object == Collection<M> \|\| WV
	private static final Logger LOG = LoggerFactory.getLogger(WindowOperatorImpl.class);

	private final WindowOperatorSpec<M, K, Object> windowOpSpec;
	private final Clock clock;
	private final WindowInternal<M, K, Object> window;
	private final FoldLeftFunction<M, Object> foldLeftFn;
	private final SupplierFunction<Object> initializer;
	private final MapFunction<M, K> keyFn;

	private final TriggerScheduler<K> triggerScheduler;
	private final Map<TriggerKey<K>, TriggerImplHandler> triggers = new ConcurrentHashMap<>();
	private TimeSeriesStore<K, Object> timeSeriesStore;

	public WindowOperatorImpl(WindowOperatorSpec<M, K, Object> windowOpSpec, Clock clock) {
	this.windowOpSpec = windowOpSpec;
	this.clock = clock;
	this.window = windowOpSpec.getWindow();
	this.foldLeftFn = window.getFoldLeftFunction();
	this.initializer = window.getInitializer();
	this.keyFn = window.getKeyExtractor();
	this.triggerScheduler= new TriggerScheduler(clock);
	}

	@Override
	protected void handleInit(Context context) {

	KeyValueStore<TimeSeriesKey<K>, Object> store =
	(KeyValueStore<TimeSeriesKey<K>, Object>) context.getTaskContext().getStore(windowOpSpec.getOpId());

	if (initializer != null) {
	initializer.init(context);
	}

	if (keyFn != null) {
	keyFn.init(context);
	}

	// For aggregating windows, we use the store in over-write mode since we only retain the aggregated
	// value. Else, we use the store in append-mode.
	if (foldLeftFn != null) {
	foldLeftFn.init(context);
	timeSeriesStore = new TimeSeriesStoreImpl(store, false);
	} else {
	timeSeriesStore = new TimeSeriesStoreImpl(store, true);
	}
	}

	@Override
	protected CompletionStage<Collection<WindowPane<K, Object>>> handleMessageAsync(M message, MessageCollector collector,
	TaskCoordinator coordinator) {
	LOG.trace("Processing message envelope: {}", message);
	List<WindowPane<K, Object>> results = new ArrayList<>();

	K key = (keyFn != null) ? keyFn.apply(message) : null;
	long timestamp = getWindowTimestamp(message);

	// For aggregating windows, we only store the aggregated window value.
	// For non-aggregating windows, we store all messages in the window.
	if (foldLeftFn == null) {
	timeSeriesStore.put(key, message, timestamp); // store is in append mode
	} else {
	List<Object> existingState = getValues(key, timestamp);

	Preconditions.checkState(existingState.size() <= 1, String.format("WindowState for aggregating windows " +
	"must not contain more than one entry per window. Current size: %s", existingState.size()));

	if (existingState.size() == 0) {
	LOG.trace("No existing state found for key {} Invoking initializer.", key);
	}
	Object oldVal = existingState.size() == 0 ? initializer.get() : existingState.get(0);
	Object aggregatedValue = foldLeftFn.apply(message, oldVal);

	timeSeriesStore.put(key, aggregatedValue, timestamp); // store is in over-write mode
	}

	if (window.getEarlyTrigger() != null) {
	TriggerKey<K> triggerKey = new TriggerKey<>(FiringType.EARLY, key, timestamp);
	TriggerImplHandler triggerImplHandler = getOrCreateTriggerImplHandler(triggerKey, window.getEarlyTrigger());
	Optional<WindowPane<K, Object>> maybeTriggeredPane =
	triggerImplHandler.onMessage(triggerKey, message, collector, coordinator);
	maybeTriggeredPane.ifPresent(results::add);
	}

	if (window.getDefaultTrigger() != null) {
	TriggerKey<K> triggerKey = new TriggerKey<>(FiringType.DEFAULT, key, timestamp);
	TriggerImplHandler triggerImplHandler = getOrCreateTriggerImplHandler(triggerKey, window.getDefaultTrigger());
	Optional<WindowPane<K, Object>> maybeTriggeredPane =
	triggerImplHandler.onMessage(triggerKey, message, collector, coordinator);
	maybeTriggeredPane.ifPresent(results::add);
	}

	return CompletableFuture.completedFuture(results);
	}

	@Override
	public Collection<WindowPane<K, Object>> handleTimer(MessageCollector collector, TaskCoordinator coordinator) {
	LOG.trace("Processing time triggers");
	List<WindowPane<K, Object>> results = new ArrayList<>();
	List<TriggerKey<K>> keys = triggerScheduler.runPendingCallbacks();

	for (TriggerKey<K> key : keys) {
	TriggerImplHandler triggerImplHandler = triggers.get(key);
	if (triggerImplHandler != null) {
	Optional<WindowPane<K, Object>> maybeTriggeredPane = triggerImplHandler.onTimer(key, collector, coordinator);
	maybeTriggeredPane.ifPresent(results::add);
	}
	}
	LOG.trace("Triggered panes: " + results.size());
	return results;
	}

	@Override
	protected OperatorSpec<M, WindowPane<K, Object>> getOperatorSpec() {
	return windowOpSpec;
	}

	@Override
	protected Collection<WindowPane<K, Object>> handleEndOfStream(MessageCollector collector, TaskCoordinator coordinator) {
	List<WindowPane<K, Object>> results = new ArrayList<>();
	Set<TriggerKey<K>> triggerKeys = new HashSet<>(triggers.keySet());
	for(TriggerKey<K> triggerKey : triggerKeys) {
	Optional<WindowPane<K, Object>> triggerResult = onTriggerFired(triggerKey, collector, coordinator);
	triggerResult.ifPresent(results::add);
	}
	return results;
	}

	@Override
	protected void handleClose() {
	if (foldLeftFn != null) {
	foldLeftFn.close();
	}
	if (timeSeriesStore != null) {
	timeSeriesStore.close();
	}
	if (initializer != null) {
	initializer.close();
	}
	if (keyFn != null) {
	keyFn.close();
	}
	}

	private TriggerImplHandler getOrCreateTriggerImplHandler(TriggerKey<K> triggerKey, Trigger<M> trigger) {
	TriggerImplHandler wrapper = triggers.get(triggerKey);
	if (wrapper != null) {
	LOG.trace("Returning existing trigger wrapper for {}", triggerKey);
	return wrapper;
	}

	LOG.trace("Creating a new trigger wrapper for {}", triggerKey);

	TriggerImpl<M, K> triggerImpl = TriggerImpls.createTriggerImpl(trigger, clock, triggerKey);
	wrapper = new TriggerImplHandler(triggerKey, triggerImpl);
	triggers.put(triggerKey, wrapper);

	return wrapper;
	}

	/**
	* Handles trigger firings and returns the optional result.
	*/
	private Optional<WindowPane<K, Object>> onTriggerFired(TriggerKey<K> triggerKey, MessageCollector collector,
	TaskCoordinator coordinator) {
	LOG.trace("Trigger key {} fired." , triggerKey);

	TriggerImplHandler wrapper = triggers.get(triggerKey);
	long timestamp = triggerKey.getTimestamp();
	K key = triggerKey.getKey();
	List<Object> existingState = getValues(key, timestamp);

	if (existingState == null \|\| existingState.size() == 0) {
	LOG.trace("No state found for triggerKey: {}", triggerKey);
	return Optional.empty();
	}

	Object windowVal = window.getFoldLeftFunction() == null ? existingState : existingState.get(0);

	WindowPane<K, Object> paneOutput = computePaneOutput(triggerKey, windowVal);

	// Handle different accumulation modes.
	if (window.getAccumulationMode() == AccumulationMode.DISCARDING) {
	LOG.trace("Clearing state for trigger key: {}", triggerKey);
	timeSeriesStore.remove(key, timestamp);
	}

	// Cancel all early triggers too when the default trigger fires. Also, clean all state for the key.
	// note: We don't handle late arrivals yet, So, all arrivals are either early or on-time.
	if (triggerKey.getType() == FiringType.DEFAULT) {

	LOG.trace("Default trigger fired. Canceling triggers for {}", triggerKey);

	cancelTrigger(triggerKey, true);
	cancelTrigger(new TriggerKey(FiringType.EARLY, triggerKey.getKey(), triggerKey.getTimestamp()), true);
	timeSeriesStore.remove(key, timestamp);
	}

	// Cancel non-repeating early triggers. All early triggers should be removed from the "triggers" map only after the
	// firing of their corresponding default trigger. Removing them pre-maturely (immediately after cancellation) will
	// will create a new {@link TriggerImplWrapper} instance at a future invocation of getOrCreateTriggerWrapper().
	// This would cause an already canceled trigger to fire again for the window.

	if (triggerKey.getType() == FiringType.EARLY && !wrapper.isRepeating()) {
	cancelTrigger(triggerKey, false);
	}

	return Optional.of(paneOutput);
	}

	/**
	* Computes the pane output corresponding to a {@link TriggerKey} that fired.
	*/
	private WindowPane<K, Object> computePaneOutput(TriggerKey<K> triggerKey, Object windowVal) {
	WindowKey<K> windowKey = new WindowKey(triggerKey.getKey(), Long.toString(triggerKey.getTimestamp()));
	WindowPane<K, Object> paneOutput =
	new WindowPane<>(windowKey, windowVal, window.getAccumulationMode(), triggerKey.getType());
	LOG.trace("Emitting pane output for trigger key {}", triggerKey);
	return paneOutput;
	}

	/**
	* Cancels the firing of the {@link TriggerImpl} identified by this {@link TriggerKey} and optionally removes it.
	*/
	private void cancelTrigger(TriggerKey<K> triggerKey, boolean shouldRemove) {
	TriggerImplHandler triggerImplHandler = triggers.get(triggerKey);
	if (triggerImplHandler != null) {
	triggerImplHandler.cancel();
	}
	if (shouldRemove && triggerKey != null) {
	triggers.remove(triggerKey);
	}
	}

	/**
	* Computes the timestamp of the window this message should belong to.
	*
	* In the case of tumbling windows, timestamp of a window is defined as the start timestamp of its corresponding window
	* interval. For instance, if the tumbling interval is 10 seconds, all messages that arrive between [1000, 1010]
	* are assigned to the window with timestamp "1000"
	*
	* In the case of session windows, timestamp is defined as the timestamp of the earliest message in the window.
	* For instance, if the session gap is 10 seconds, and the first message in the window arrives at "1002" seconds,
	* all messages (that arrive within 10 seconds of their previous message) are assigned a timestamp "1002".
	*
	* @param message the input message
	* @return the timestamp of the window this message should belong to
	*/
	private long getWindowTimestamp(M message) {
	if (window.getWindowType() == WindowType.TUMBLING) {
	long triggerDurationMs = ((TimeTrigger<M>) window.getDefaultTrigger()).getDuration().toMillis();
	final long now = clock.currentTimeMillis();
	// assign timestamp to be the start timestamp of the window boundary
	long timestamp = now - now % triggerDurationMs;
	return timestamp;
	} else {
	K key = keyFn.apply(message);
	// get the value with the earliest timestamp for the provided key.
	ClosableIterator<TimestampedValue<Object>> iterator = timeSeriesStore.get(key, 0, Long.MAX_VALUE, 1);
	List<TimestampedValue<Object>> timestampedValues = toList(iterator);

	// If there are no existing sessions for the key, we return the current timestamp. If not, return the
	// timestamp of the earliest message.
	long timestamp = (timestampedValues.isEmpty())? clock.currentTimeMillis() : timestampedValues.get(0).getTimestamp();

	return timestamp;
	}
	}

	/**
	* Return a list of values in the store for the provided key and timestamp
	*
	* @param key the key to look up in the store
	* @param timestamp the timestamp to look up in the store
	* @return the list of values for the provided key
	*/
	private List<Object> getValues(K key, long timestamp) {
	ClosableIterator<TimestampedValue<Object>> iterator = timeSeriesStore.get(key, timestamp);
	List<TimestampedValue<Object>> timestampedValues = toList(iterator);
	List<Object> values = timestampedValues.stream().map(element -> element.getValue()).collect(Collectors.toList());

	LOG.trace("Returning {} for key {} and timestamp {}", new Object[] {values, key, timestamp});
	return values;
	}

	/**
	* Returns an unmodifiable list of all elements in the provided iterator.
	* The iterator is guaranteed to be closed after its execution.
	*
	* @param iterator the provided iterator.
	* @param <V> the type of elements in the iterator
	* @return a list of all elements returned by the iterator
	*/
	static <V> List<V> toList(ClosableIterator<V> iterator) {
	Preconditions.checkNotNull(iterator);

	List<V> values = new ArrayList<>();
	try {
	while (iterator.hasNext()) {
	values.add(iterator.next());
	}
	} finally {
	iterator.close();
	}
	return Collections.unmodifiableList(values);
	}

	/**
	* State corresponding to a created {@link TriggerImpl} instance.
	*/
	private class TriggerImplHandler {
	// The context, and the {@link TriggerImpl} instance corresponding to this triggerKey
	private final TriggerImpl<M, K> impl;
	// Guard to ensure that we don't invoke onMessage or onTimer on already cancelled triggers
	private boolean isCancelled = false;

	public TriggerImplHandler(TriggerKey<K> key, TriggerImpl<M, K> impl) {
	this.impl = impl;
	}

	public Optional<WindowPane<K, Object>> onMessage(TriggerKey<K> triggerKey, M message,
	MessageCollector collector, TaskCoordinator coordinator) {
	if (!isCancelled) {
	LOG.trace("Forwarding callbacks for {}", message);
	impl.onMessage(message, triggerScheduler);

	if (impl.shouldFire()) {
	// repeating trigger can trigger multiple times, So, clear the state to allow future triggerings.
	if (impl instanceof RepeatingTriggerImpl) {
	((RepeatingTriggerImpl<M, K>) impl).clear();
	}
	return onTriggerFired(triggerKey, collector, coordinator);
	}
	}
	return Optional.empty();
	}

	public Optional<WindowPane<K, Object>> onTimer(TriggerKey<K> key, MessageCollector collector,
	TaskCoordinator coordinator) {
	if (impl.shouldFire() && !isCancelled) {
	LOG.trace("Triggering timer triggers");

	// repeating trigger can trigger multiple times, So, clear the trigger to allow future triggerings.
	if (impl instanceof RepeatingTriggerImpl) {
	((RepeatingTriggerImpl<M, K>) impl).clear();
	}
	return onTriggerFired(key, collector, coordinator);
	}
	return Optional.empty();
	}

	public void cancel() {
	impl.cancel();
	isCancelled = true;
	}

	public boolean isRepeating() {
	return this.impl instanceof RepeatingTriggerImpl;
	}
	}
	}