/**
 * 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.pulsar.broker.service;

import org.apache.pulsar.common.policies.data.Policies;
import org.apache.pulsar.common.policies.data.PublishRate;

import java.util.concurrent.atomic.LongAdder;

public class PublishRateLimiterImpl implements PublishRateLimiter {
    protected volatile int publishMaxMessageRate = 0;
    protected volatile long publishMaxByteRate = 0;
    protected volatile boolean publishThrottlingEnabled = false;
    protected volatile boolean publishRateExceeded = false;
    protected volatile LongAdder currentPublishMsgCount = new LongAdder();
    protected volatile LongAdder currentPublishByteCount = new LongAdder();

    public PublishRateLimiterImpl(Policies policies, String clusterName) {
        update(policies, clusterName);
    }

    public PublishRateLimiterImpl(PublishRate maxPublishRate) {
        update(maxPublishRate);
    }

    @Override
    public void checkPublishRate() {
        if (this.publishThrottlingEnabled && !publishRateExceeded) {
            long currentPublishMsgRate = this.currentPublishMsgCount.sum();
            long currentPublishByteRate = this.currentPublishByteCount.sum();
            if ((this.publishMaxMessageRate > 0 && currentPublishMsgRate > this.publishMaxMessageRate)
                    || (this.publishMaxByteRate > 0 && currentPublishByteRate > this.publishMaxByteRate)) {
                publishRateExceeded = true;
            }
        }
    }

    @Override
    public void incrementPublishCount(int numOfMessages, long msgSizeInBytes) {
        if (this.publishThrottlingEnabled) {
            this.currentPublishMsgCount.add(numOfMessages);
            this.currentPublishByteCount.add(msgSizeInBytes);
        }
    }

    @Override
    public boolean resetPublishCount() {
        if (this.publishThrottlingEnabled) {
            this.currentPublishMsgCount.reset();
            this.currentPublishByteCount.reset();
            this.publishRateExceeded = false;
            return true;
        }
        return false;
    }

    @Override
    public boolean isPublishRateExceeded() {
        return publishRateExceeded;
    }

    @Override
    public void update(Policies policies, String clusterName) {
        final PublishRate maxPublishRate = policies.publishMaxMessageRate != null
                ? policies.publishMaxMessageRate.get(clusterName)
                : null;
        update(maxPublishRate);
    }

    public void update(PublishRate maxPublishRate) {
        if (maxPublishRate != null
            && (maxPublishRate.publishThrottlingRateInMsg > 0 || maxPublishRate.publishThrottlingRateInByte > 0)) {
            this.publishThrottlingEnabled = true;
            this.publishMaxMessageRate = Math.max(maxPublishRate.publishThrottlingRateInMsg, 0);
            this.publishMaxByteRate = Math.max(maxPublishRate.publishThrottlingRateInByte, 0);
        } else {
            this.publishMaxMessageRate = 0;
            this.publishMaxByteRate = 0;
            this.publishThrottlingEnabled = false;
        }
        resetPublishCount();
    }

    @Override
    public boolean tryAcquire(int numbers, long bytes) {
        return false;
    }
}

@@ -1799,24 +1799,75 @@
 Exceeded) %7B%0A
+            if (this.publishMaxByteRate %3E 0) %7B%0A    
             
@@ -1881,27 +1881,28 @@
 rrentPublish
-Msg
+Byte
 Rate = this.
@@ -1911,27 +1911,28 @@
 rrentPublish
-Msg
+Byte
 Count.sum();
@@ -1936,37 +1936,40 @@
 ();%0A            
-long 
+    if (
 currentPublishBy
@@ -1979,45 +1979,145 @@
 ate 
-=
+%3E
  this.
-currentPublishByteCount.sum();
+publishMaxByteRate) %7B%0A                    publishRateExceeded = true;%0A                    return;%0A                %7D%0A            %7D%0A
 %0A   
@@ -2128,17 +2128,16 @@
      if 
-(
 (this.pu
@@ -2151,35 +2151,56 @@
 xMessageRate %3E 0
- &&
+) %7B%0A                long
  currentPublishM
@@ -2210,93 +2210,65 @@
 ate 
-%3E
+=
  this.
-publishMaxMessageRate)%0A                    %7C%7C (this.publishMaxByteRate %3E 0 && 
+currentPublishMsgCount.sum();%0A                if (
 curr
@@ -2269,36 +2269,35 @@
  (currentPublish
-Byte
+Msg
 Rate %3E this.publ
@@ -2306,18 +2306,20 @@
 hMax
-Byt
+Messag
 eRate)
-)
  %7B%0A 
@@ -2325,32 +2325,36 @@
                 
+ 
+   
 publishRateExcee
@@ -2357,32 +2357,50 @@
 xceeded = true;%0A
+                %7D%0A
             %7D%0A  

improve performance of checkPublishRate() (#8623)