Skip to content

ZOOKEEPER-5074: Quota checks ignore earlier operations in the same multi transaction #2432

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SahilRajputdotexe wants to merge 1 commit into
base: master
Choose a base branch
from
+155 −11
Open

ZOOKEEPER-5074: Quota checks ignore earlier operations in the same multi transaction #2432

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions zookeeper-server/src/main/java/org/apache/zookeeper/server/PrepRequestProcessor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import org.apache.zookeeper.KeeperException.Code;
import org.apache.zookeeper.MultiOperationRecord;
import org.apache.zookeeper.Op;
import org.apache.zookeeper.StatsTrack;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooDefs.OpCode;
import org.apache.zookeeper.common.ConfigException;
Expand Down Expand Up @@ -104,6 +105,17 @@ public class PrepRequestProcessor extends ZooKeeperCriticalThread implements Req
private final boolean digestEnabled;
private DigestCalculator digestCalculator;

/**
* The quota-relevant changes (node count and data bytes, per quota
* prefix) of the operations that were already validated as part of the
* request currently being processed. The operations of a multi
* transaction are validated one by one before any of them is applied to
* the data tree, so the quota stat nodes do not yet account for the
* earlier operations of the same transaction. Only accessed from the
* request processing thread.
*/
private final Map<String, StatsTrack> pendingQuotaChanges = new HashMap<>();

ZooKeeperServer zks;

public enum DigestOpCode {
Expand Down Expand Up @@ -361,6 +373,7 @@ protected void pRequest2Txn(int type, long zxid, Request request, Record record)
if (nodeRecord.childCount > 0) {
throw new KeeperException.NotEmptyException(path);
}
zks.checkQuota(path, nodeRecord.data, null, OpCode.delete, pendingQuotaChanges);
request.setTxn(new DeleteTxn(path));
parentRecord = parentRecord.duplicate(request.getHdr().getZxid());
parentRecord.childCount--;
Expand All @@ -381,7 +394,7 @@ protected void pRequest2Txn(int type, long zxid, Request request, Record record)
validatePath(path, request.sessionId);
nodeRecord = getRecordForPath(path);
zks.checkACL(request.cnxn, nodeRecord.acl, ZooDefs.Perms.WRITE, request.authInfo, path, null);
zks.checkQuota(path, nodeRecord.data, setDataRequest.getData(), OpCode.setData);
zks.checkQuota(path, nodeRecord.data, setDataRequest.getData(), OpCode.setData, pendingQuotaChanges);
int newVersion = checkAndIncVersion(nodeRecord.stat.getVersion(), setDataRequest.getVersion(), path);
request.setTxn(new SetDataTxn(path, setDataRequest.getData(), newVersion));
nodeRecord = nodeRecord.duplicate(request.getHdr().getZxid());
Expand Down Expand Up @@ -679,7 +692,7 @@ private void pRequest2TxnCreate(int type, Request request, Record record) throws
throw new KeeperException.NoChildrenForEphemeralsException(path);
}
int newCversion = parentRecord.stat.getCversion() + 1;
zks.checkQuota(path, null, data, OpCode.create);
zks.checkQuota(path, null, data, OpCode.create, pendingQuotaChanges);
if (type == OpCode.createContainer) {
request.setTxn(new CreateContainerTxn(path, data, listACL, newCversion));
} else if (type == OpCode.createTTL) {
Expand Down Expand Up @@ -758,6 +771,7 @@ private static int checkAndIncVersion(int currentVersion, int expectedVersion, S
protected void pRequest(Request request) throws RequestProcessorException {
request.setHdr(null);
request.setTxn(null);
pendingQuotaChanges.clear();

if (!request.isThrottled()) {
pRequestHelper(request);
Expand Down
64 changes: 55 additions & 9 deletions zookeeper-server/src/main/java/org/apache/zookeeper/server/ZooKeeperServer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2111,30 +2111,64 @@ public void checkACL(ServerCnxn cnxn, List<ACL> acl, int perm, List<Id> ids, Str
* @param data
* the data to be set, or {@code null} for none
* @param type
* currently, create and setData need to check quota
* currently, create and setData need to check quota; delete
* only records the released count and bytes
* @param pendingChanges
* accumulates, per quota prefix, the count and byte changes
* made by the operations of the current request that were
* already validated; a multi transaction is validated before
* any of it is applied, so the stat nodes alone do not
* reflect the earlier operations of the transaction
*/
public void checkQuota(String path, byte[] lastData, byte[] data, int type) throws KeeperException.QuotaExceededException {
public void checkQuota(String path, byte[] lastData, byte[] data, int type,
Map<String, StatsTrack> pendingChanges) throws KeeperException.QuotaExceededException {
if (!enforceQuota) {
return;
}
long dataBytes = (data == null) ? 0 : data.length;
long lastDataBytes = (lastData == null) ? 0 : lastData.length;
ZKDatabase zkDatabase = getZKDatabase();
String lastPrefix = zkDatabase.getDataTree().getMaxPrefixWithQuota(path);
if (StringUtils.isEmpty(lastPrefix)) {
return;
}

final String namespace = PathUtils.getTopNamespace(path);
long bytesDiff;
long countDiff;
switch (type) {
case OpCode.create:
checkQuota(lastPrefix, dataBytes, 1, namespace);
bytesDiff = dataBytes;
countDiff = 1;
break;
case OpCode.setData:
checkQuota(lastPrefix, dataBytes - (lastData == null ? 0 : lastData.length), 0, namespace);
bytesDiff = dataBytes - lastDataBytes;
countDiff = 0;
break;
case OpCode.delete:
// a delete cannot exceed a quota, but the released count and
// bytes have to be remembered so that later operations in the
// same transaction are checked against the correct usage
bytesDiff = -lastDataBytes;
countDiff = -1;
break;
default:
throw new IllegalArgumentException("Unsupported OpCode for checkQuota: " + type);
}

if (type != OpCode.delete) {
final String namespace = PathUtils.getTopNamespace(path);
checkQuota(lastPrefix, bytesDiff, countDiff, namespace, pendingChanges);
}

StatsTrack pending = pendingChanges.get(lastPrefix);
if (pending == null) {
pending = new StatsTrack();
pending.setCount(0);
pending.setBytes(0);
pendingChanges.put(lastPrefix, pending);
}
pending.setCount(pending.getCount() + countDiff);
pending.setBytes(pending.getBytes() + bytesDiff);
}

/**
Expand All @@ -2148,9 +2182,13 @@ public void checkQuota(String path, byte[] lastData, byte[] data, int type) thro
* the diff to be added to the count
* @param namespace
* the namespace for collecting quota exceeded errors
* @param pendingChanges
* the count and byte changes made by operations that were
* already validated as part of the current request but are
* not yet reflected in the quota stat node
*/
private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, String namespace)
throws KeeperException.QuotaExceededException {
private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, String namespace,
Map<String, StatsTrack> pendingChanges) throws KeeperException.QuotaExceededException {
LOG.debug("checkQuota: lastPrefix={}, bytesDiff={}, countDiff={}", lastPrefix, bytesDiff, countDiff);

// now check the quota we set
Expand Down Expand Up @@ -2187,9 +2225,17 @@ private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, Strin
currentStats = new StatsTrack(node.data);
}

long pendingCount = 0;
long pendingBytes = 0;
StatsTrack pending = pendingChanges.get(lastPrefix);
if (pending != null) {
pendingCount = pending.getCount();
pendingBytes = pending.getBytes();
}

//check the Count Quota
if (checkCountQuota) {
long newCount = currentStats.getCount() + countDiff;
long newCount = currentStats.getCount() + pendingCount + countDiff;
boolean isCountHardLimit = limitStats.getCountHardLimit() > -1;
long countLimit = isCountHardLimit ? limitStats.getCountHardLimit() : limitStats.getCount();

Expand All @@ -2205,7 +2251,7 @@ private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, Strin

//check the Byte Quota
if (checkByteQuota) {
long newBytes = currentStats.getBytes() + bytesDiff;
long newBytes = currentStats.getBytes() + pendingBytes + bytesDiff;
boolean isByteHardLimit = limitStats.getByteHardLimit() > -1;
long byteLimit = isByteHardLimit ? limitStats.getByteHardLimit() : limitStats.getBytes();
if (newBytes > byteLimit) {
Expand Down
84 changes: 84 additions & 0 deletions zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -473,6 +473,90 @@ public void testMultiCreateThenSetDataShouldFail() throws Exception {
assertNull(zk.exists(subPath, null));
}

@Test
public void testMultiCreatesExceedingCountHardQuotaShouldFail() throws Exception {
final String path = "/c";

zk.create(path, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);

final StatsTrack st = new StatsTrack();
st.setCountHardLimit(3);
SetQuotaCommand.createQuota(zk, path, st);

// each create on its own stays within the limit; only the
// transaction as a whole exceeds it
final List<Op> ops = Arrays.asList(
Op.create(path + "/1", null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT),
Op.create(path + "/2", null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT),
Op.create(path + "/3", null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT));

try {
zk.multi(ops);
fail("should fail transaction when the creates together exceed the count hard quota");
} catch (QuotaExceededException e) {
//expected
}

assertNull(zk.exists(path + "/1", null));
assertNull(zk.exists(path + "/2", null));
assertNull(zk.exists(path + "/3", null));
}

@Test
public void testMultiCreatesExceedingBytesHardQuotaShouldFail() throws Exception {
final String path = "/b";

zk.create(path, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);

final byte[] data5b = "Hello".getBytes(StandardCharsets.UTF_8);

final StatsTrack st = new StatsTrack();
st.setByteHardLimit(3 * data5b.length - 1);
SetQuotaCommand.createQuota(zk, path, st);

final List<Op> ops = Arrays.asList(
Op.create(path + "/1", data5b, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT),
Op.create(path + "/2", data5b, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT),
Op.create(path + "/3", data5b, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT));

try {
zk.multi(ops);
fail("should fail transaction when the creates together exceed the byte hard quota");
} catch (QuotaExceededException e) {
//expected
}

assertNull(zk.exists(path + "/1", null));
assertNull(zk.exists(path + "/2", null));
assertNull(zk.exists(path + "/3", null));
}

@Test
public void testMultiDeleteThenCreateWithinCountHardQuotaShouldWork() throws Exception {
final String path = "/replace";
final String oldChild = path + "/old";
final String newChild = path + "/new";

zk.create(path, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);

final StatsTrack st = new StatsTrack();
st.setCountHardLimit(2);
SetQuotaCommand.createQuota(zk, path, st);

zk.create(oldChild, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);

// the node count is at the hard limit, but the transaction deletes
// as much as it creates, so it should be allowed
final List<Op> ops = Arrays.asList(
Op.delete(oldChild, -1),
Op.create(newChild, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT));

zk.multi(ops);

assertNull(zk.exists(oldChild, null));
assertNotNull(zk.exists(newChild, null));
}

@Test
public void testDeleteBytesQuota() throws Exception {

Expand Down