geode-redis/src/main/java/org/apache/geode/redis/internal/pubsub/PublishResultCollector.java

/*
 * 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.geode.redis.internal.pubsub;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.geode.redis.internal.netty.Client;


public class PublishResultCollector {
  private final CountDownLatch countDownLatch;
  private final Subscriptions subscriptions;
  private final AtomicLong successCount = new AtomicLong();

  public PublishResultCollector(int resultsExpected, Subscriptions subscriptions) {
    countDownLatch = new CountDownLatch(resultsExpected);
    this.subscriptions = subscriptions;
  }

  public void success() {
    successCount.incrementAndGet();
    countDownLatch.countDown();
  }

  public void failure(Client result) {
    pruneFailure(result);
    countDownLatch.countDown();
  }

  public long getSuccessCount() {
    waitForResults();
    return successCount.get();
  }

  private void pruneFailure(Client c) {
    if (c.isDead()) {
      subscriptions.remove(c);
    }
  }

  private void waitForResults() {
    try {
      countDownLatch.await();
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
  }
}