ratis-server/src/main/java/org/apache/ratis/server/impl/PendingRequests.java - ratis - Git at Google

 /*
  * 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.ratis.server.impl;

 import org.apache.ratis.conf.RaftProperties;
 import org.apache.ratis.proto.RaftProtos.CommitInfoProto;
 import org.apache.ratis.protocol.Message;
 import org.apache.ratis.protocol.exceptions.NotLeaderException;
 import org.apache.ratis.protocol.RaftClientReply;
 import org.apache.ratis.protocol.RaftClientRequest;
 import org.apache.ratis.protocol.exceptions.RaftException;
 import org.apache.ratis.protocol.RaftGroupMemberId;
 import org.apache.ratis.protocol.SetConfigurationRequest;
 import org.apache.ratis.server.RaftServerConfigKeys;
 import org.apache.ratis.server.metrics.RaftServerMetricsImpl;
 import org.apache.ratis.statemachine.TransactionContext;
 import org.apache.ratis.util.JavaUtils;
 import org.apache.ratis.util.Preconditions;
 import org.apache.ratis.util.ResourceSemaphore;
 import org.apache.ratis.util.SizeInBytes;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;

 class PendingRequests {
   public static final Logger LOG = LoggerFactory.getLogger(PendingRequests.class);

   private static final int ONE_MB = SizeInBytes.ONE_MB.getSizeInt();

   /**
    * Round up to the nearest MB.
    */
   static int roundUpMb(long bytes) {
     return Math.toIntExact((bytes - 1) / ONE_MB + 1);
   }

   static class Permit {}

   static class RequestLimits extends ResourceSemaphore.Group {
     RequestLimits(int elementLimit, int megabyteLimit) {
       super(elementLimit, megabyteLimit);
     }

     int getElementCount() {
       return get(0).used();
     }

     int getMegaByteSize() {
       return get(1).used();
     }

     ResourceSemaphore.ResourceAcquireStatus tryAcquire(int messageSizeMb) {
       return tryAcquire(1, messageSizeMb);
     }

     void releaseExtraMb(int extraMb) {
       release(0, extraMb);
     }

     void release(int diffMb) {
       release(1, diffMb);
     }
   }

   private static class RequestMap {
     private final Object name;
     private final ConcurrentMap<Long, PendingRequest> map = new ConcurrentHashMap<>();
     private final RaftServerMetricsImpl raftServerMetrics;

     /** Permits to put new requests, always synchronized. */
     private final Map<Permit, Permit> permits = new HashMap<>();
     /** Track and limit the number of requests and the total message size. */
     private final RequestLimits resource;
     /** The size (in byte) of all the requests in this map. */
     private final AtomicLong requestSize = new AtomicLong();


     RequestMap(Object name, int elementLimit, int megabyteLimit, RaftServerMetricsImpl raftServerMetrics) {
       this.name = name;
       this.resource = new RequestLimits(elementLimit, megabyteLimit);
       this.raftServerMetrics = raftServerMetrics;

       raftServerMetrics.addNumPendingRequestsGauge(resource::getElementCount);
       raftServerMetrics.addNumPendingRequestsMegaByteSize(resource::getMegaByteSize);
     }

     Permit tryAcquire(Message message) {
       final int messageSize = Message.getSize(message);
       final int messageSizeMb = roundUpMb(messageSize );
       final ResourceSemaphore.ResourceAcquireStatus acquired = resource.tryAcquire(messageSizeMb);
       LOG.trace("tryAcquire {} MB? {}", messageSizeMb, acquired);
       if (acquired == ResourceSemaphore.ResourceAcquireStatus.FAILED_IN_ELEMENT_LIMIT) {
         raftServerMetrics.onRequestQueueLimitHit();
         raftServerMetrics.onResourceLimitHit();
         return null;
       } else if (acquired == ResourceSemaphore.ResourceAcquireStatus.FAILED_IN_BYTE_SIZE_LIMIT) {
         raftServerMetrics.onRequestByteSizeLimitHit();
         raftServerMetrics.onResourceLimitHit();
         return null;
       }

       // release extra MB
       final long oldSize = requestSize.getAndAdd(messageSize);
       final long newSize = oldSize + messageSize;
       final int diffMb = roundUpMb(newSize) - roundUpMb(oldSize);
       if (messageSizeMb > diffMb) {
         resource.releaseExtraMb(messageSizeMb - diffMb);
       }
       return putPermit();
     }

     private synchronized Permit putPermit() {
       if (resource.isClosed()) {
         return null;
       }
       final Permit permit = new Permit();
       permits.put(permit, permit);
       return permit;
     }

     synchronized PendingRequest put(Permit permit, long index, PendingRequest p) {
       LOG.debug("{}: PendingRequests.put {} -> {}", name, index, p);
       final Permit removed = permits.remove(permit);
       if (removed == null) {
         return null;
       }
       Preconditions.assertTrue(removed == permit);
       final PendingRequest previous = map.put(index, p);
       Preconditions.assertTrue(previous == null);
       return p;
     }

     PendingRequest get(long index) {
       final PendingRequest r = map.get(index);
       LOG.debug("{}: PendingRequests.get {} returns {}", name, index, r);
       return r;
     }

     PendingRequest remove(long index) {
       final PendingRequest r = map.remove(index);
       LOG.debug("{}: PendingRequests.remove {} returns {}", name, index, r);
       if (r == null) {
         return null;
       }
       final int messageSize = Message.getSize(r.getRequest().getMessage());
       final long oldSize = requestSize.getAndAdd(-messageSize);
       final long newSize = oldSize - messageSize;
       final int diffMb = roundUpMb(oldSize) - roundUpMb(newSize);
       resource.release(diffMb);
       LOG.trace("release {} MB", diffMb);
       return r;
     }

     Collection<TransactionContext> setNotLeaderException(NotLeaderException nle,
                                                          Collection<CommitInfoProto> commitInfos) {
       synchronized (this) {
         resource.close();
         permits.clear();
       }

       LOG.debug("{}: PendingRequests.setNotLeaderException", name);
       final List<TransactionContext> transactions = new ArrayList<>(map.size());
       for(;;) {
         final Iterator<Long> i = map.keySet().iterator();
         if (!i.hasNext()) { // the map is empty
           return transactions;
         }

         final PendingRequest pending = map.remove(i.next());
         if (pending != null) {
           transactions.add(pending.setNotLeaderException(nle, commitInfos));
         }
       }
     }

     void close() {
       if (raftServerMetrics != null) {
         raftServerMetrics.removeNumPendingRequestsGauge();
         raftServerMetrics.removeNumPendingRequestsByteSize();
       }
     }
   }

   private PendingRequest pendingSetConf;
   private final String name;
   private final RequestMap pendingRequests;

   PendingRequests(RaftGroupMemberId id, RaftProperties properties, RaftServerMetricsImpl raftServerMetrics) {
     this.name = id + "-" + JavaUtils.getClassSimpleName(getClass());
     this.pendingRequests = new RequestMap(id,
         RaftServerConfigKeys.Write.elementLimit(properties),
         Math.toIntExact(
             RaftServerConfigKeys.Write.byteLimit(properties).getSize()
                 / SizeInBytes.ONE_MB.getSize()), //round down
         raftServerMetrics);
   }

   Permit tryAcquire(Message message) {
     return pendingRequests.tryAcquire(message);
   }

   PendingRequest add(Permit permit, RaftClientRequest request, TransactionContext entry) {
     // externally synced for now
     final long index = entry.getLogEntry().getIndex();
     LOG.debug("{}: addPendingRequest at index={}, request={}", name, index, request);
     final PendingRequest pending = new PendingRequest(index, request, entry);
     return pendingRequests.put(permit, index, pending);
   }

   PendingRequest addConfRequest(SetConfigurationRequest request) {
     Preconditions.assertTrue(pendingSetConf == null);
     pendingSetConf = new PendingRequest(request);
     return pendingSetConf;
   }

   void replySetConfiguration(Function<RaftClientRequest, RaftClientReply> newSuccessReply) {
     // we allow the pendingRequest to be null in case that the new leader
     // commits the new configuration while it has not received the retry
     // request from the client
     if (pendingSetConf != null) {
       final RaftClientRequest request = pendingSetConf.getRequest();
       LOG.debug("{}: sends success for {}", name, request);
       // for setConfiguration we do not need to wait for statemachine. send back
       // reply after it's committed.
       pendingSetConf.setReply(newSuccessReply.apply(request));
       pendingSetConf = null;
     }
   }

   void failSetConfiguration(RaftException e) {
     Preconditions.assertTrue(pendingSetConf != null);
     pendingSetConf.setException(e);
     pendingSetConf = null;
   }

   TransactionContext getTransactionContext(long index) {
     PendingRequest pendingRequest = pendingRequests.get(index);
     // it is possible that the pendingRequest is null if this peer just becomes
     // the new leader and commits transactions received by the previous leader
     return pendingRequest != null ? pendingRequest.getEntry() : null;
   }

   void replyPendingRequest(long index, RaftClientReply reply) {
     final PendingRequest pending = pendingRequests.remove(index);
     if (pending != null) {
       Preconditions.assertTrue(pending.getIndex() == index);
       pending.setReply(reply);
     }
   }

   /**
    * The leader state is stopped. Send NotLeaderException to all the pending
    * requests since they have not got applied to the state machine yet.
    */
   Collection<TransactionContext> sendNotLeaderResponses(NotLeaderException nle,
                                                         Collection<CommitInfoProto> commitInfos) {
     LOG.info("{}: sendNotLeaderResponses", name);

     final Collection<TransactionContext> transactions = pendingRequests.setNotLeaderException(nle, commitInfos);
     if (pendingSetConf != null) {
       pendingSetConf.setNotLeaderException(nle, commitInfos);
     }
     return transactions;
   }

   void close() {
     if (pendingRequests != null) {
       pendingRequests.close();
     }
   }
 }
	/*
	* 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.ratis.server.impl;

	import org.apache.ratis.conf.RaftProperties;
	import org.apache.ratis.proto.RaftProtos.CommitInfoProto;
	import org.apache.ratis.protocol.Message;
	import org.apache.ratis.protocol.exceptions.NotLeaderException;
	import org.apache.ratis.protocol.RaftClientReply;
	import org.apache.ratis.protocol.RaftClientRequest;
	import org.apache.ratis.protocol.exceptions.RaftException;
	import org.apache.ratis.protocol.RaftGroupMemberId;
	import org.apache.ratis.protocol.SetConfigurationRequest;
	import org.apache.ratis.server.RaftServerConfigKeys;
	import org.apache.ratis.server.metrics.RaftServerMetricsImpl;
	import org.apache.ratis.statemachine.TransactionContext;
	import org.apache.ratis.util.JavaUtils;
	import org.apache.ratis.util.Preconditions;
	import org.apache.ratis.util.ResourceSemaphore;
	import org.apache.ratis.util.SizeInBytes;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import java.util.concurrent.ConcurrentHashMap;
	import java.util.concurrent.ConcurrentMap;
	import java.util.concurrent.atomic.AtomicLong;
	import java.util.function.Function;

	class PendingRequests {
	public static final Logger LOG = LoggerFactory.getLogger(PendingRequests.class);

	private static final int ONE_MB = SizeInBytes.ONE_MB.getSizeInt();

	/**
	* Round up to the nearest MB.
	*/
	static int roundUpMb(long bytes) {
	return Math.toIntExact((bytes - 1) / ONE_MB + 1);
	}

	static class Permit {}

	static class RequestLimits extends ResourceSemaphore.Group {
	RequestLimits(int elementLimit, int megabyteLimit) {
	super(elementLimit, megabyteLimit);
	}

	int getElementCount() {
	return get(0).used();
	}

	int getMegaByteSize() {
	return get(1).used();
	}

	ResourceSemaphore.ResourceAcquireStatus tryAcquire(int messageSizeMb) {
	return tryAcquire(1, messageSizeMb);
	}

	void releaseExtraMb(int extraMb) {
	release(0, extraMb);
	}

	void release(int diffMb) {
	release(1, diffMb);
	}
	}

	private static class RequestMap {
	private final Object name;
	private final ConcurrentMap<Long, PendingRequest> map = new ConcurrentHashMap<>();
	private final RaftServerMetricsImpl raftServerMetrics;

	/** Permits to put new requests, always synchronized. */
	private final Map<Permit, Permit> permits = new HashMap<>();
	/** Track and limit the number of requests and the total message size. */
	private final RequestLimits resource;
	/** The size (in byte) of all the requests in this map. */
	private final AtomicLong requestSize = new AtomicLong();


	RequestMap(Object name, int elementLimit, int megabyteLimit, RaftServerMetricsImpl raftServerMetrics) {
	this.name = name;
	this.resource = new RequestLimits(elementLimit, megabyteLimit);
	this.raftServerMetrics = raftServerMetrics;

	raftServerMetrics.addNumPendingRequestsGauge(resource::getElementCount);
	raftServerMetrics.addNumPendingRequestsMegaByteSize(resource::getMegaByteSize);
	}

	Permit tryAcquire(Message message) {
	final int messageSize = Message.getSize(message);
	final int messageSizeMb = roundUpMb(messageSize );
	final ResourceSemaphore.ResourceAcquireStatus acquired = resource.tryAcquire(messageSizeMb);
	LOG.trace("tryAcquire {} MB? {}", messageSizeMb, acquired);
	if (acquired == ResourceSemaphore.ResourceAcquireStatus.FAILED_IN_ELEMENT_LIMIT) {
	raftServerMetrics.onRequestQueueLimitHit();
	raftServerMetrics.onResourceLimitHit();
	return null;
	} else if (acquired == ResourceSemaphore.ResourceAcquireStatus.FAILED_IN_BYTE_SIZE_LIMIT) {
	raftServerMetrics.onRequestByteSizeLimitHit();
	raftServerMetrics.onResourceLimitHit();
	return null;
	}

	// release extra MB
	final long oldSize = requestSize.getAndAdd(messageSize);
	final long newSize = oldSize + messageSize;
	final int diffMb = roundUpMb(newSize) - roundUpMb(oldSize);
	if (messageSizeMb > diffMb) {
	resource.releaseExtraMb(messageSizeMb - diffMb);
	}
	return putPermit();
	}

	private synchronized Permit putPermit() {
	if (resource.isClosed()) {
	return null;
	}
	final Permit permit = new Permit();
	permits.put(permit, permit);
	return permit;
	}

	synchronized PendingRequest put(Permit permit, long index, PendingRequest p) {
	LOG.debug("{}: PendingRequests.put {} -> {}", name, index, p);
	final Permit removed = permits.remove(permit);
	if (removed == null) {
	return null;
	}
	Preconditions.assertTrue(removed == permit);
	final PendingRequest previous = map.put(index, p);
	Preconditions.assertTrue(previous == null);
	return p;
	}

	PendingRequest get(long index) {
	final PendingRequest r = map.get(index);
	LOG.debug("{}: PendingRequests.get {} returns {}", name, index, r);
	return r;
	}

	PendingRequest remove(long index) {
	final PendingRequest r = map.remove(index);
	LOG.debug("{}: PendingRequests.remove {} returns {}", name, index, r);
	if (r == null) {
	return null;
	}
	final int messageSize = Message.getSize(r.getRequest().getMessage());
	final long oldSize = requestSize.getAndAdd(-messageSize);
	final long newSize = oldSize - messageSize;
	final int diffMb = roundUpMb(oldSize) - roundUpMb(newSize);
	resource.release(diffMb);
	LOG.trace("release {} MB", diffMb);
	return r;
	}

	Collection<TransactionContext> setNotLeaderException(NotLeaderException nle,
	Collection<CommitInfoProto> commitInfos) {
	synchronized (this) {
	resource.close();
	permits.clear();
	}

	LOG.debug("{}: PendingRequests.setNotLeaderException", name);
	final List<TransactionContext> transactions = new ArrayList<>(map.size());
	for(;;) {
	final Iterator<Long> i = map.keySet().iterator();
	if (!i.hasNext()) { // the map is empty
	return transactions;
	}

	final PendingRequest pending = map.remove(i.next());
	if (pending != null) {
	transactions.add(pending.setNotLeaderException(nle, commitInfos));
	}
	}
	}

	void close() {
	if (raftServerMetrics != null) {
	raftServerMetrics.removeNumPendingRequestsGauge();
	raftServerMetrics.removeNumPendingRequestsByteSize();
	}
	}
	}

	private PendingRequest pendingSetConf;
	private final String name;
	private final RequestMap pendingRequests;

	PendingRequests(RaftGroupMemberId id, RaftProperties properties, RaftServerMetricsImpl raftServerMetrics) {
	this.name = id + "-" + JavaUtils.getClassSimpleName(getClass());
	this.pendingRequests = new RequestMap(id,
	RaftServerConfigKeys.Write.elementLimit(properties),
	Math.toIntExact(
	RaftServerConfigKeys.Write.byteLimit(properties).getSize()
	/ SizeInBytes.ONE_MB.getSize()), //round down
	raftServerMetrics);
	}

	Permit tryAcquire(Message message) {
	return pendingRequests.tryAcquire(message);
	}

	PendingRequest add(Permit permit, RaftClientRequest request, TransactionContext entry) {
	// externally synced for now
	final long index = entry.getLogEntry().getIndex();
	LOG.debug("{}: addPendingRequest at index={}, request={}", name, index, request);
	final PendingRequest pending = new PendingRequest(index, request, entry);
	return pendingRequests.put(permit, index, pending);
	}

	PendingRequest addConfRequest(SetConfigurationRequest request) {
	Preconditions.assertTrue(pendingSetConf == null);
	pendingSetConf = new PendingRequest(request);
	return pendingSetConf;
	}

	void replySetConfiguration(Function<RaftClientRequest, RaftClientReply> newSuccessReply) {
	// we allow the pendingRequest to be null in case that the new leader
	// commits the new configuration while it has not received the retry
	// request from the client
	if (pendingSetConf != null) {
	final RaftClientRequest request = pendingSetConf.getRequest();
	LOG.debug("{}: sends success for {}", name, request);
	// for setConfiguration we do not need to wait for statemachine. send back
	// reply after it's committed.
	pendingSetConf.setReply(newSuccessReply.apply(request));
	pendingSetConf = null;
	}
	}

	void failSetConfiguration(RaftException e) {
	Preconditions.assertTrue(pendingSetConf != null);
	pendingSetConf.setException(e);
	pendingSetConf = null;
	}

	TransactionContext getTransactionContext(long index) {
	PendingRequest pendingRequest = pendingRequests.get(index);
	// it is possible that the pendingRequest is null if this peer just becomes
	// the new leader and commits transactions received by the previous leader
	return pendingRequest != null ? pendingRequest.getEntry() : null;
	}

	void replyPendingRequest(long index, RaftClientReply reply) {
	final PendingRequest pending = pendingRequests.remove(index);
	if (pending != null) {
	Preconditions.assertTrue(pending.getIndex() == index);
	pending.setReply(reply);
	}
	}

	/**
	* The leader state is stopped. Send NotLeaderException to all the pending
	* requests since they have not got applied to the state machine yet.
	*/
	Collection<TransactionContext> sendNotLeaderResponses(NotLeaderException nle,
	Collection<CommitInfoProto> commitInfos) {
	LOG.info("{}: sendNotLeaderResponses", name);

	final Collection<TransactionContext> transactions = pendingRequests.setNotLeaderException(nle, commitInfos);
	if (pendingSetConf != null) {
	pendingSetConf.setNotLeaderException(nle, commitInfos);
	}
	return transactions;
	}

	void close() {
	if (pendingRequests != null) {
	pendingRequests.close();
	}
	}
	}