[ISSUE #3949] add test cases

2026-09-24 16:04:00 +08:00 · 2022-07-13 11:29:39 +08:00
parent aa313e96f8
commit cc827021c4
9 changed files with 299 additions and 18 deletions
@@ -48,6 +48,13 @@ public class CustomizedRetryPolicy implements RetryPolicy {
        TimeUnit.HOURS.toMillis(2)
    };

+    public CustomizedRetryPolicy() {
+    }
+
+    public CustomizedRetryPolicy(long[] next) {
+        this.next = next;
+    }
+
    public long[] getNext() {
        return next;
    }
@@ -28,6 +28,15 @@ public class ExponentialRetryPolicy implements RetryPolicy {
    private long max = TimeUnit.HOURS.toMillis(2);
    private long multiplier = 2;

+    public ExponentialRetryPolicy() {
+    }
+
+    public ExponentialRetryPolicy(long initial, long max, long multiplier) {
+        this.initial = initial;
+        this.max = max;
+        this.multiplier = multiplier;
+    }
+
    public long getInitial() {
        return initial;
    }
@@ -24,7 +24,6 @@ import com.google.common.collect.ComparisonChain;
 import io.grpc.Context;
 import io.grpc.stub.StreamObserver;
 import io.netty.channel.ChannelId;
-import java.net.SocketAddress;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicReference;
 import org.apache.rocketmq.common.message.MessageExt;
@@ -39,7 +38,6 @@ import org.apache.rocketmq.proxy.service.relay.ProxyChannel;
 import org.apache.rocketmq.proxy.service.relay.ProxyRelayResult;
 import org.apache.rocketmq.proxy.service.relay.ProxyRelayService;
 import org.apache.rocketmq.proxy.service.transaction.TransactionId;
-import org.apache.rocketmq.remoting.common.RemotingUtil;
 import org.apache.rocketmq.remoting.protocol.RemotingCommand;

 public class GrpcClientChannel extends ProxyChannel {
@@ -117,16 +115,6 @@ public class GrpcClientChannel extends ProxyChannel {
        return this.telemetryCommandRef.get() != null;
    }

-    @Override
-    protected SocketAddress localAddress0() {
-        return RemotingUtil.string2SocketAddress(this.localAddress);
-    }
-
-    @Override
-    protected SocketAddress remoteAddress0() {
-        return RemotingUtil.string2SocketAddress(this.remoteAddress);
-    }
-
    @Override
    protected CompletableFuture<Void> processOtherMessage(Object msg) {
        if (msg instanceof TelemetryCommand) {
@@ -153,7 +153,7 @@ public class ClientActivity extends AbstractMessingActivity {
            ProxyContext context = createContext(ctx);
            String clientId = context.getVal(GrpcContextConstants.CLIENT_ID);
            LanguageCode languageCode = context.getVal(GrpcContextConstants.LANGUAGE);
-            Settings clientSettings = grpcClientSettingsManager.removeClientSettings(clientId);
+            Settings clientSettings = grpcClientSettingsManager.removeAndGetClientSettings(context);

            switch (clientSettings.getClientType()) {
                case PRODUCER:
@@ -75,6 +75,9 @@ public class GrpcClientSettingsManager {
    public Settings getClientSettings(ProxyContext ctx) {
        String clientId = ctx.getVal(GrpcContextConstants.CLIENT_ID);
        Settings settings = CLIENT_SETTINGS_MAP.get(clientId);
+        if (settings == null) {
+            return null;
+        }
        if (settings.hasSubscription()) {
            settings = mergeSubscriptionData(ctx, settings,
                GrpcConverter.wrapResourceWithNamespace(settings.getSubscription().getGroup()));
@@ -141,7 +144,17 @@ public class GrpcClientSettingsManager {
        CLIENT_SETTINGS_MAP.put(clientId, settings);
    }

-    public Settings removeClientSettings(String clientId) {
-        return CLIENT_SETTINGS_MAP.remove(clientId);
+    public void removeClientSettings(String clientId) {
+        CLIENT_SETTINGS_MAP.remove(clientId);
+    }
+
+    public Settings removeAndGetClientSettings(ProxyContext ctx) {
+        String clientId = ctx.getVal(GrpcContextConstants.CLIENT_ID);
+        Settings settings = CLIENT_SETTINGS_MAP.remove(clientId);
+        if (settings == null) {
+            return null;
+        }
+        return mergeSubscriptionData(ctx, settings,
+            GrpcConverter.wrapResourceWithNamespace(settings.getSubscription().getGroup()));
    }
 }
@@ -44,12 +44,15 @@ import org.apache.rocketmq.logging.InternalLoggerFactory;
 import org.apache.rocketmq.proxy.common.ContextVariable;
 import org.apache.rocketmq.proxy.common.ProxyContext;
 import org.apache.rocketmq.proxy.service.transaction.TransactionId;
+import org.apache.rocketmq.remoting.common.RemotingUtil;
 import org.apache.rocketmq.remoting.protocol.RemotingCommand;

 public abstract class ProxyChannel extends AbstractChannel {
    private static final InternalLogger log = InternalLoggerFactory.getLogger(LoggerName.PROXY_LOGGER_NAME);
    protected final String remoteAddress;
+    protected final SocketAddress remoteSocketAddress;
    protected final String localAddress;
+    protected final SocketAddress localSocketAddress;

    protected final ProxyRelayService proxyRelayService;

@@ -57,14 +60,18 @@ public abstract class ProxyChannel extends AbstractChannel {
        super(parent);
        this.proxyRelayService = proxyRelayService;
        this.remoteAddress = remoteAddress;
+        this.remoteSocketAddress = RemotingUtil.string2SocketAddress(remoteAddress);
        this.localAddress = localAddress;
+        this.localSocketAddress = RemotingUtil.string2SocketAddress(localAddress);
    }

    protected ProxyChannel(ProxyRelayService proxyRelayService, Channel parent, ChannelId id, String remoteAddress, String localAddress) {
        super(parent, id);
        this.proxyRelayService = proxyRelayService;
        this.remoteAddress = remoteAddress;
+        this.remoteSocketAddress = RemotingUtil.string2SocketAddress(remoteAddress);
        this.localAddress = localAddress;
+        this.localSocketAddress = RemotingUtil.string2SocketAddress(localAddress);
    }

    @Override
@@ -179,4 +186,14 @@ public abstract class ProxyChannel extends AbstractChannel {
    protected void doWrite(ChannelOutboundBuffer in) throws Exception {

    }
+
+    @Override
+    protected SocketAddress localAddress0() {
+        return this.localSocketAddress;
+    }
+
+    @Override
+    protected SocketAddress remoteAddress0() {
+        return this.remoteSocketAddress;
+    }
 }
@@ -49,7 +49,6 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.when;

@@ -180,7 +179,7 @@ public class ClientActivityTest extends BaseActivityTest {
    public void testProducerNotifyClientTermination() throws Throwable {
        Context context = createContext();

-        when(this.grpcClientSettingsManager.removeClientSettings(eq(CLIENT_ID))).thenReturn(Settings.newBuilder()
+        when(this.grpcClientSettingsManager.removeAndGetClientSettings(any())).thenReturn(Settings.newBuilder()
            .setClientType(ClientType.PRODUCER)
            .setPublishing(Publishing.newBuilder()
                .addTopics(Resource.newBuilder().setName(TOPIC).build())
@@ -207,7 +206,7 @@ public class ClientActivityTest extends BaseActivityTest {
    public void testConsumerNotifyClientTermination() throws Throwable {
        Context context = createContext();

-        when(this.grpcClientSettingsManager.removeClientSettings(eq(CLIENT_ID))).thenReturn(Settings.newBuilder()
+        when(this.grpcClientSettingsManager.removeAndGetClientSettings(any())).thenReturn(Settings.newBuilder()
            .setClientType(ClientType.PUSH_CONSUMER)
            .build());
        ArgumentCaptor<ClientChannelInfo> channelInfoArgumentCaptor = ArgumentCaptor.forClass(ClientChannelInfo.class);
@@ -0,0 +1,98 @@
+/*
+ * 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.rocketmq.proxy.grpc.v2.common;
+
+import apache.rocketmq.v2.CustomizedBackoff;
+import apache.rocketmq.v2.ExponentialBackoff;
+import apache.rocketmq.v2.Resource;
+import apache.rocketmq.v2.RetryPolicy;
+import apache.rocketmq.v2.Settings;
+import apache.rocketmq.v2.Subscription;
+import com.google.protobuf.util.Durations;
+import org.apache.rocketmq.common.subscription.CustomizedRetryPolicy;
+import org.apache.rocketmq.common.subscription.ExponentialRetryPolicy;
+import org.apache.rocketmq.common.subscription.GroupRetryPolicyType;
+import org.apache.rocketmq.common.subscription.SubscriptionGroupConfig;
+import org.apache.rocketmq.proxy.common.ProxyContext;
+import org.apache.rocketmq.proxy.grpc.v2.BaseActivityTest;
+import org.apache.rocketmq.proxy.grpc.v2.GrpcContextConstants;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.when;
+
+public class GrpcClientSettingsManagerTest extends BaseActivityTest {
+    private GrpcClientSettingsManager grpcClientSettingsManager;
+
+    @Before
+    public void before() throws Throwable {
+        super.before();
+        this.grpcClientSettingsManager = new GrpcClientSettingsManager(this.messagingProcessor);
+    }
+
+    @Test
+    public void testGetSubscriptionData() {
+        SubscriptionGroupConfig subscriptionGroupConfig = new SubscriptionGroupConfig();
+        when(this.messagingProcessor.getSubscriptionGroupConfig(any(), any()))
+            .thenReturn(subscriptionGroupConfig);
+
+        this.grpcClientSettingsManager.updateClientSettings(CLIENT_ID, Settings.newBuilder()
+            .setSubscription(Subscription.newBuilder()
+                .setGroup(Resource.newBuilder().setName("group").build())
+                .build())
+            .build());
+
+        ProxyContext context = ProxyContext.create().withVal(GrpcContextConstants.CLIENT_ID, CLIENT_ID);
+
+        Settings settings = this.grpcClientSettingsManager.getClientSettings(context);
+        assertEquals(settings.getBackoffPolicy(), GrpcClientSettingsManager.DEFAULT_CONSUMER_SETTINGS.getBackoffPolicy());
+
+        subscriptionGroupConfig.setRetryMaxTimes(3);
+        subscriptionGroupConfig.getGroupRetryPolicy().setType(GroupRetryPolicyType.CUSTOMIZED);
+        subscriptionGroupConfig.getGroupRetryPolicy().setCustomizedRetryPolicy(new CustomizedRetryPolicy(new long[]{1000}));
+        settings = this.grpcClientSettingsManager.getClientSettings(context);
+        assertEquals(RetryPolicy.newBuilder()
+            .setMaxAttempts(3)
+            .setCustomizedBackoff(CustomizedBackoff.newBuilder()
+                .addNext(Durations.fromSeconds(1))
+                .build())
+            .build(), settings.getBackoffPolicy());
+
+        subscriptionGroupConfig.setRetryMaxTimes(10);
+        subscriptionGroupConfig.getGroupRetryPolicy().setType(GroupRetryPolicyType.EXPONENTIAL);
+        subscriptionGroupConfig.getGroupRetryPolicy().setExponentialRetryPolicy(new ExponentialRetryPolicy(1000, 2000, 3));
+        settings = this.grpcClientSettingsManager.getClientSettings(context);
+        assertEquals(RetryPolicy.newBuilder()
+            .setMaxAttempts(10)
+            .setExponentialBackoff(ExponentialBackoff.newBuilder()
+                .setMax(Durations.fromSeconds(2))
+                .setInitial(Durations.fromSeconds(1))
+                .setMultiplier(3)
+                .build())
+            .build(), settings.getBackoffPolicy());
+
+        Settings settings1 = this.grpcClientSettingsManager.removeAndGetClientSettings(context);
+        assertEquals(settings, settings1);
+
+        assertNull(this.grpcClientSettingsManager.getClientSettings(context));
+        assertNull(this.grpcClientSettingsManager.removeAndGetClientSettings(context));
+    }
+}
@@ -0,0 +1,150 @@
+/*
+ * 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.rocketmq.proxy.service.relay;
+
+import io.netty.channel.Channel;
+import java.net.SocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.util.UUID;
+import java.util.concurrent.CompletableFuture;
+import org.apache.rocketmq.common.message.MessageClientIDSetter;
+import org.apache.rocketmq.common.message.MessageDecoder;
+import org.apache.rocketmq.common.message.MessageExt;
+import org.apache.rocketmq.common.protocol.RequestCode;
+import org.apache.rocketmq.common.protocol.body.ConsumeMessageDirectlyResult;
+import org.apache.rocketmq.common.protocol.body.ConsumerRunningInfo;
+import org.apache.rocketmq.common.protocol.header.CheckTransactionStateRequestHeader;
+import org.apache.rocketmq.common.protocol.header.ConsumeMessageDirectlyResultRequestHeader;
+import org.apache.rocketmq.common.protocol.header.GetConsumerRunningInfoRequestHeader;
+import org.apache.rocketmq.proxy.service.transaction.TransactionId;
+import org.apache.rocketmq.remoting.common.RemotingUtil;
+import org.apache.rocketmq.remoting.protocol.RemotingCommand;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import static org.junit.Assert.*;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.when;
+
+@RunWith(MockitoJUnitRunner.class)
+public class ProxyChannelTest {
+
+    @Mock
+    private ProxyRelayService proxyRelayService;
+
+    protected abstract static class MockProxyChannel extends ProxyChannel {
+
+        protected MockProxyChannel(ProxyRelayService proxyRelayService, Channel parent,
+            String remoteAddress, String localAddress) {
+            super(proxyRelayService, parent, remoteAddress, localAddress);
+        }
+
+        @Override public boolean isOpen() {
+            return false;
+        }
+
+        @Override public boolean isActive() {
+            return false;
+        }
+    }
+
+    @Test
+    public void testWriteAndFlush() throws Exception {
+        ArgumentCaptor<ConsumeMessageDirectlyResultRequestHeader> consumeMessageDirectlyArgumentCaptor =
+            ArgumentCaptor.forClass(ConsumeMessageDirectlyResultRequestHeader.class);
+        when(this.proxyRelayService.processConsumeMessageDirectly(any(), any(), consumeMessageDirectlyArgumentCaptor.capture()))
+            .thenReturn(new CompletableFuture<>());
+
+        ArgumentCaptor<GetConsumerRunningInfoRequestHeader> getConsumerRunningInfoArgumentCaptor =
+            ArgumentCaptor.forClass(GetConsumerRunningInfoRequestHeader.class);
+        when(this.proxyRelayService.processGetConsumerRunningInfo(any(), any(), getConsumerRunningInfoArgumentCaptor.capture()))
+            .thenReturn(new CompletableFuture<>());
+
+        CheckTransactionStateRequestHeader checkTransactionStateRequestHeader = new CheckTransactionStateRequestHeader();
+        checkTransactionStateRequestHeader.setBrokerName("broker");
+        checkTransactionStateRequestHeader.setTransactionId(MessageClientIDSetter.createUniqID());
+        RemotingCommand checkTransactionRequest = RemotingCommand.createRequestCommand(RequestCode.CHECK_TRANSACTION_STATE, checkTransactionStateRequestHeader);
+        MessageExt transactionMessageExt = new MessageExt();
+        transactionMessageExt.setTopic("topic");
+        transactionMessageExt.setTags("tags");
+        transactionMessageExt.setBornHost(RemotingUtil.string2SocketAddress("127.0.0.2:8888"));
+        transactionMessageExt.setStoreHost(RemotingUtil.string2SocketAddress("127.0.0.1:10911"));
+        transactionMessageExt.setBody(UUID.randomUUID().toString().getBytes(StandardCharsets.UTF_8));
+        transactionMessageExt.setMsgId(MessageClientIDSetter.createUniqID());
+        checkTransactionRequest.setBody(MessageDecoder.encode(transactionMessageExt, false));
+
+        GetConsumerRunningInfoRequestHeader consumerRunningInfoRequestHeader = new GetConsumerRunningInfoRequestHeader();
+        consumerRunningInfoRequestHeader.setConsumerGroup("group");
+        consumerRunningInfoRequestHeader.setClientId("clientId");
+        RemotingCommand consumerRunningInfoRequest = RemotingCommand.createRequestCommand(RequestCode.GET_CONSUMER_RUNNING_INFO, consumerRunningInfoRequestHeader);
+
+        ConsumeMessageDirectlyResultRequestHeader consumeMessageDirectlyResultRequestHeader = new ConsumeMessageDirectlyResultRequestHeader();
+        consumeMessageDirectlyResultRequestHeader.setConsumerGroup("group");
+        consumeMessageDirectlyResultRequestHeader.setClientId("clientId");
+        MessageExt consumeMessageDirectlyMessageExt = new MessageExt();
+        consumeMessageDirectlyMessageExt.setTopic("topic");
+        consumeMessageDirectlyMessageExt.setTags("tags");
+        consumeMessageDirectlyMessageExt.setBornHost(RemotingUtil.string2SocketAddress("127.0.0.2:8888"));
+        consumeMessageDirectlyMessageExt.setStoreHost(RemotingUtil.string2SocketAddress("127.0.0.1:10911"));
+        consumeMessageDirectlyMessageExt.setBody(UUID.randomUUID().toString().getBytes(StandardCharsets.UTF_8));
+        consumeMessageDirectlyMessageExt.setMsgId(MessageClientIDSetter.createUniqID());
+        RemotingCommand consumeMessageDirectlyResult = RemotingCommand.createRequestCommand(RequestCode.CONSUME_MESSAGE_DIRECTLY, consumeMessageDirectlyResultRequestHeader);
+        consumeMessageDirectlyResult.setBody(MessageDecoder.encode(consumeMessageDirectlyMessageExt, false));
+
+        MockProxyChannel channel = new MockProxyChannel(this.proxyRelayService, null, "127.0.0.2:8888", "127.0.0.1:10911") {
+            @Override
+            protected CompletableFuture<Void> processOtherMessage(Object msg) {
+                return CompletableFuture.completedFuture(null);
+            }
+
+            @Override
+            protected CompletableFuture<Void> processCheckTransaction(CheckTransactionStateRequestHeader header,
+                MessageExt messageExt, TransactionId transactionId) {
+                assertEquals(checkTransactionStateRequestHeader, header);
+                assertArrayEquals(transactionMessageExt.getBody(), messageExt.getBody());
+                return CompletableFuture.completedFuture(null);
+            }
+
+            @Override
+            protected CompletableFuture<Void> processGetConsumerRunningInfo(RemotingCommand command,
+                GetConsumerRunningInfoRequestHeader header,
+                CompletableFuture<ProxyRelayResult<ConsumerRunningInfo>> responseFuture) {
+                assertEquals(consumerRunningInfoRequestHeader, getConsumerRunningInfoArgumentCaptor.getValue());
+                assertEquals(consumerRunningInfoRequestHeader, header);
+                return CompletableFuture.completedFuture(null);
+            }
+
+            @Override
+            protected CompletableFuture<Void> processConsumeMessageDirectly(RemotingCommand command,
+                ConsumeMessageDirectlyResultRequestHeader header, MessageExt messageExt,
+                CompletableFuture<ProxyRelayResult<ConsumeMessageDirectlyResult>> responseFuture) {
+                assertEquals(consumeMessageDirectlyResultRequestHeader, consumeMessageDirectlyArgumentCaptor.getValue());
+                assertEquals(consumeMessageDirectlyResultRequestHeader, header);
+                assertArrayEquals(consumeMessageDirectlyMessageExt.getBody(), messageExt.getBody());
+                return CompletableFuture.completedFuture(null);
+            }
+        };
+
+        assertTrue(channel.writeAndFlush(checkTransactionRequest).isSuccess());
+        assertTrue(channel.writeAndFlush(consumerRunningInfoRequest).isSuccess());
+        assertTrue(channel.writeAndFlush(consumeMessageDirectlyResult).isSuccess());
+    }
+}