[pinpoint-apm#11346] Optimize agent lookup for agent compatibility

emeroad · Aug 13, 2024 · 9ddc23a · 9ddc23a
1 parent ccd7e28
commit 9ddc23a
Show file tree

Hide file tree

Showing 7 changed files with 89 additions and 23 deletions.
diff --git a/collector/src/main/java/com/navercorp/pinpoint/collector/dao/AgentInfoDao.java b/collector/src/main/java/com/navercorp/pinpoint/collector/dao/AgentInfoDao.java
@@ -25,5 +25,5 @@
 public interface AgentInfoDao {
     void insert(AgentInfoBo agentInfo);
 
-    AgentInfoBo getAgentInfo(String agentId, long timestamp);
+    AgentInfoBo getSimpleAgentInfo(String agentId, long timestamp);
 }
diff --git a/collector/src/main/java/com/navercorp/pinpoint/collector/dao/hbase/HbaseAgentInfoDao.java b/collector/src/main/java/com/navercorp/pinpoint/collector/dao/hbase/HbaseAgentInfoDao.java
@@ -21,9 +21,11 @@
 import com.navercorp.pinpoint.common.hbase.HbaseColumnFamily;
 import com.navercorp.pinpoint.common.hbase.HbaseOperations;
 import com.navercorp.pinpoint.common.hbase.ResultsExtractor;
+import com.navercorp.pinpoint.common.hbase.RowMapper;
 import com.navercorp.pinpoint.common.hbase.TableNameProvider;
 import com.navercorp.pinpoint.common.server.bo.AgentInfoBo;
 import com.navercorp.pinpoint.common.server.bo.serializer.agent.AgentIdRowKeyEncoder;
+import com.navercorp.pinpoint.common.server.dao.hbase.mapper.SingleResultsExtractor;
 import com.navercorp.pinpoint.common.server.util.RowKeyUtils;
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.Put;
@@ -53,10 +55,11 @@ public class HbaseAgentInfoDao implements AgentInfoDao {
 
     public HbaseAgentInfoDao(HbaseOperations hbaseTemplate,
                              TableNameProvider tableNameProvider,
-                             ResultsExtractor<AgentInfoBo> agentInfoResultsExtractor) {
+                             RowMapper<AgentInfoBo> agentInfoMapper) {
         this.hbaseTemplate = Objects.requireNonNull(hbaseTemplate, "hbaseTemplate");
         this.tableNameProvider = Objects.requireNonNull(tableNameProvider, "tableNameProvider");
-        this.agentInfoResultsExtractor = Objects.requireNonNull(agentInfoResultsExtractor, "agentInfoResultsExtractor");
+
+        this.agentInfoResultsExtractor = new SingleResultsExtractor<>(agentInfoMapper);
     }
 
     @Override
@@ -94,11 +97,23 @@ public void insert(AgentInfoBo agentInfo) {
         hbaseTemplate.put(agentInfoTableName, put);
     }
 
-    public AgentInfoBo getAgentInfo(final String agentId, final long timestamp) {
+    public AgentInfoBo getSimpleAgentInfo(final String agentId, final long timestamp) {
         Objects.requireNonNull(agentId, "agentId");
 
-        final Scan scan = createScan(agentId, timestamp);
         final TableName agentInfoTableName = tableNameProvider.getTableName(DESCRIPTOR.getTable());
+        return getSimpleAgentInfoBoByScanner(agentId, timestamp, agentInfoTableName);
+//        return getAgentInfoBoByGet(agentId, timestamp, agentInfoTableName);
+    }
+
+//    private AgentInfoBo getSimpleAgentInfoBoByGet(String agentId, long timestamp, TableName agentInfoTableName) {
+//        byte[] rowKey = rowKeyEncoder.encodeRowKey(agentId, timestamp);
+//        Get get = new Get(rowKey);
+//        get.addColumn(DESCRIPTOR.getName(), DESCRIPTOR.QUALIFIER_IDENTIFIER);
+//        return hbaseTemplate.get(agentInfoTableName, get, agentInfoMapper);
+//    }
+
+    private AgentInfoBo getSimpleAgentInfoBoByScanner(String agentId, long timestamp, TableName agentInfoTableName) {
+        final Scan scan = createScan(agentId, timestamp);
         return this.hbaseTemplate.find(agentInfoTableName, scan, agentInfoResultsExtractor);
     }
 
@@ -110,10 +125,13 @@ private Scan createScan(String agentId, long currentTime) {
 
         scan.withStartRow(startKeyBytes);
         scan.withStopRow(endKeyBytes);
-        scan.addFamily(DESCRIPTOR.getName());
+
         scan.readVersions(1);
+        scan.setOneRowLimit();
         scan.setCaching(SCANNER_CACHING);
 
+        scan.addColumn(DESCRIPTOR.getName(), DESCRIPTOR.QUALIFIER_IDENTIFIER);
+
         return scan;
     }
 }
diff --git a/...n/java/com/navercorp/pinpoint/collector/receiver/grpc/service/AgentLifecycleListener.java b/...n/java/com/navercorp/pinpoint/collector/receiver/grpc/service/AgentLifecycleListener.java
@@ -20,11 +20,12 @@
 import com.navercorp.pinpoint.collector.service.AgentInfoService;
 import com.navercorp.pinpoint.collector.util.ManagedAgentLifeCycle;
 import com.navercorp.pinpoint.common.server.bo.AgentInfoBo;
+import com.navercorp.pinpoint.common.trace.ServiceType;
 import com.navercorp.pinpoint.grpc.Header;
-import com.navercorp.pinpoint.grpc.server.lifecycle.PingSession;
 import com.navercorp.pinpoint.grpc.server.lifecycle.LifecycleListener;
-import org.apache.logging.log4j.Logger;
+import com.navercorp.pinpoint.grpc.server.lifecycle.PingSession;
 import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
 import java.util.Objects;
 
@@ -51,12 +52,16 @@ public void connect(PingSession lifecycle) {
         logger.info("connect:{}", lifecycle);
         final Header header = lifecycle.getHeader();
         try {
-            final AgentInfoBo agentInfoBo = agentInfoService.getAgentInfo(header.getAgentId(), header.getAgentStartTime());
-            if (null != agentInfoBo) {
-                lifecycle.setServiceType(agentInfoBo.getServiceTypeCode());
+            if (lifecycle.getServiceType() == ServiceType.UNDEFINED.getCode()) {
+                // fallback
+                final AgentInfoBo agentInfoBo = agentInfoService.getSimpleAgentInfo(header.getAgentId(), header.getAgentStartTime());
+                logger.info("ServiceType is UNDEFINED. Fallback:AgentInfo lookup {} -> {}", lifecycle, agentInfoBo);
+                if (agentInfoBo != null) {
+                    lifecycle.setServiceType(agentInfoBo.getServiceTypeCode());
+                }
             }
         } catch (Exception e) {
-            logger.warn("Failed to handle. ping session={}", lifecycle, e);
+            logger.warn("Fallback:AgentInfo lookup Failed. session={}", lifecycle, e);
         }
         lifecycleService.updateState(lifecycle, ManagedAgentLifeCycle.RUNNING);
     }

diff --git a/collector/src/main/java/com/navercorp/pinpoint/collector/service/AgentInfoService.java b/collector/src/main/java/com/navercorp/pinpoint/collector/service/AgentInfoService.java
@@ -50,7 +50,7 @@ public void insert(@Valid final AgentInfoBo agentInfoBo) {
         applicationIndexDao.insert(agentInfoBo);
     }
 
-    public AgentInfoBo getAgentInfo(@NotBlank final String agentId, @PositiveOrZero final long timestamp) {
-        return agentInfoDao.getAgentInfo(agentId, timestamp);
+    public AgentInfoBo getSimpleAgentInfo(@NotBlank final String agentId, @PositiveOrZero final long timestamp) {
+        return agentInfoDao.getSimpleAgentInfo(agentId, timestamp);
     }
 }
diff --git a/...e/mapper/AgentInfoBoResultsExtractor.java → .../hbase/mapper/SingleResultsExtractor.java b/...e/mapper/AgentInfoBoResultsExtractor.java → .../hbase/mapper/SingleResultsExtractor.java
@@ -19,7 +19,6 @@
 import com.google.common.collect.Iterables;
 import com.navercorp.pinpoint.common.hbase.ResultsExtractor;
 import com.navercorp.pinpoint.common.hbase.RowMapper;
-import com.navercorp.pinpoint.common.server.bo.AgentInfoBo;
 import org.apache.hadoop.hbase.client.Result;
 import org.apache.hadoop.hbase.client.ResultScanner;
 import org.springframework.stereotype.Component;
@@ -30,21 +29,25 @@
  * @author HyunGil Jeong
  */
 @Component
-public class AgentInfoBoResultsExtractor implements ResultsExtractor<AgentInfoBo> {
+public class SingleResultsExtractor<T> implements ResultsExtractor<T> {
 
-    private final RowMapper<AgentInfoBo> agentInfoMapper;
+    private final RowMapper<T> mapper;
 
-    public AgentInfoBoResultsExtractor(RowMapper<AgentInfoBo> agentInfoMapper) {
-        this.agentInfoMapper = Objects.requireNonNull(agentInfoMapper, "agentInfoMapper");
+    public SingleResultsExtractor(RowMapper<T> mapper) {
+        this.mapper = Objects.requireNonNull(mapper, "mapper");
+    }
+
+    public RowMapper<T> getMapper() {
+        return mapper;
     }
 
     @Override
-    public AgentInfoBo extractData(ResultScanner results) throws Exception {
+    public T extractData(ResultScanner results) throws Exception {
         final Result first = Iterables.getFirst(results, null);
         if (first == null) {
             return null;
         }
-        return agentInfoMapper.mapRow(first, 0);
+        return mapper.mapRow(first, 0);
 
     }
 }
diff --git a/grpc/src/main/java/com/navercorp/pinpoint/grpc/server/lifecycle/PingSession.java b/grpc/src/main/java/com/navercorp/pinpoint/grpc/server/lifecycle/PingSession.java
@@ -53,11 +53,16 @@ public long nextEventIdAllocator() {
     }
 
     public short getServiceType() {
-        return serviceType;
+        if (serviceType != ServiceType.UNDEFINED.getCode()) {
+            return serviceType;
+        }
+        return (short) header.getServiceType();
     }
 
     public void setServiceType(short serviceType) {
-        this.serviceType = serviceType;
+        if (header.getServiceType() == ServiceType.UNDEFINED.getCode()) {
+            this.serviceType = serviceType;
+        }
     }
 
     // Flag to avoid duplication.

diff --git a/grpc/src/test/java/com/navercorp/pinpoint/grpc/server/lifecycle/PingSessionTest.java b/grpc/src/test/java/com/navercorp/pinpoint/grpc/server/lifecycle/PingSessionTest.java
@@ -0,0 +1,35 @@
+package com.navercorp.pinpoint.grpc.server.lifecycle;
+
+import com.navercorp.pinpoint.common.trace.ServiceType;
+import com.navercorp.pinpoint.grpc.Header;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+
+class PingSessionTest {
+
+    @Test
+    void getServiceType() {
+        Header header = new Header("name", "agentId", "agentName", "appName",
+                ServiceType.SPRING.getCode(),  11, 22, Collections.emptyList());
+        PingSession session = new PingSession(1L, header);
+
+        Assertions.assertEquals(ServiceType.SPRING.getCode(), session.getServiceType());
+
+        session.setServiceType(ServiceType.TEST.getCode());
+        Assertions.assertEquals(ServiceType.SPRING.getCode(), session.getServiceType());
+    }
+
+    @Test
+    void getServiceType_undefined() {
+        Header header = new Header("name", "agentId", "agentName", "appName",
+                ServiceType.UNDEFINED.getCode(),  11, 22, Collections.emptyList());
+        PingSession session = new PingSession(1L, header);
+
+        Assertions.assertEquals(ServiceType.UNDEFINED.getCode(), session.getServiceType());
+
+        session.setServiceType(ServiceType.SPRING.getCode());
+        Assertions.assertEquals(ServiceType.SPRING.getCode(), session.getServiceType());
+    }
+}