Add trace model (#94)

Author: hanahmily

.gitignore

@@ -20,4 +20,5 @@ packages/
 /test/plugin/workspace
 /test/jacoco/classes
 /test/jacoco/*.exec
-test/jacoco
+test/jacoco
+CLAUDE.md

CHANGES.md

@@ -18,6 +18,7 @@ Release Notes.
 * Bump up parent Apache pom to v35.
 * Bump up maven to 3.6.3.
 * Add IDEA setup doc to support super large generated file(by protoc).
+* Add Trace model.
 
 0.8.0
 ------------------

src/main/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClient.java

@@ -43,6 +43,7 @@
 import org.apache.skywalking.banyandb.database.v1.BanyandbDatabase.IndexRule;
 import org.apache.skywalking.banyandb.database.v1.BanyandbDatabase.IndexRuleBinding;
 import org.apache.skywalking.banyandb.database.v1.BanyandbDatabase.Subject;
+import org.apache.skywalking.banyandb.database.v1.BanyandbDatabase.Trace;
 import org.apache.skywalking.banyandb.model.v1.BanyandbModel;
 import org.apache.skywalking.banyandb.property.v1.BanyandbProperty;
 import org.apache.skywalking.banyandb.property.v1.BanyandbProperty.Property;
@@ -53,6 +54,8 @@
 import org.apache.skywalking.banyandb.measure.v1.MeasureServiceGrpc;
 import org.apache.skywalking.banyandb.stream.v1.BanyandbStream;
 import org.apache.skywalking.banyandb.stream.v1.StreamServiceGrpc;
+import org.apache.skywalking.banyandb.trace.v1.BanyandbTrace;
+import org.apache.skywalking.banyandb.trace.v1.TraceServiceGrpc;
 import org.apache.skywalking.banyandb.v1.client.auth.AuthInterceptor;
 import org.apache.skywalking.banyandb.v1.client.grpc.HandleExceptionsWith;
 import org.apache.skywalking.banyandb.v1.client.grpc.channel.ChannelManager;
@@ -69,6 +72,7 @@
 import org.apache.skywalking.banyandb.v1.client.metadata.ResourceExist;
 import org.apache.skywalking.banyandb.v1.client.metadata.StreamMetadataRegistry;
 import org.apache.skywalking.banyandb.v1.client.metadata.TopNAggregationMetadataRegistry;
+import org.apache.skywalking.banyandb.v1.client.metadata.TraceMetadataRegistry;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -124,6 +128,11 @@ public class BanyanDBClient implements Closeable {
      */
     @Getter(value = AccessLevel.PACKAGE)
     private MeasureServiceGrpc.MeasureServiceStub measureServiceStub;
+    /**
+     * gRPC client stub
+     */
+    @Getter(value = AccessLevel.PACKAGE)
+    private TraceServiceGrpc.TraceServiceStub traceServiceStub;
     /**
      * gRPC future stub.
      */
@@ -134,6 +143,11 @@ public class BanyanDBClient implements Closeable {
      */
     @Getter(value = AccessLevel.PACKAGE)
     private MeasureServiceGrpc.MeasureServiceBlockingStub measureServiceBlockingStub;
+    /**
+     * gRPC future stub.
+     */
+    @Getter(value = AccessLevel.PACKAGE)
+    private TraceServiceGrpc.TraceServiceBlockingStub traceServiceBlockingStub;
     /**
      * The connection status.
      */
@@ -211,8 +225,10 @@ public void connect() throws IOException {
                 this.channel = interceptedChannel;
                 streamServiceBlockingStub = StreamServiceGrpc.newBlockingStub(this.channel);
                 measureServiceBlockingStub = MeasureServiceGrpc.newBlockingStub(this.channel);
+                traceServiceBlockingStub = TraceServiceGrpc.newBlockingStub(this.channel);
                 streamServiceStub = StreamServiceGrpc.newStub(this.channel);
                 measureServiceStub = MeasureServiceGrpc.newStub(this.channel);
+                traceServiceStub = TraceServiceGrpc.newStub(this.channel);
                 isConnected = true;
             }
         } finally {
@@ -228,8 +244,10 @@ void connect(Channel channel) {
                 this.channel = channel;
                 streamServiceBlockingStub = StreamServiceGrpc.newBlockingStub(this.channel);
                 measureServiceBlockingStub = MeasureServiceGrpc.newBlockingStub(this.channel);
+                traceServiceBlockingStub = TraceServiceGrpc.newBlockingStub(this.channel);
                 streamServiceStub = StreamServiceGrpc.newStub(this.channel);
                 measureServiceStub = MeasureServiceGrpc.newStub(this.channel);
+                traceServiceStub = TraceServiceGrpc.newStub(this.channel);
                 isConnected = true;
             }
         } finally {
@@ -396,6 +414,49 @@ public StreamWrite createStreamWrite(String group, String name, final String ele
         return new StreamWrite(this.metadataCache.findStreamMetadata(group, name), elementId, timestamp);
     }
 
+    /**
+     * Build a trace bulk write processor.
+     *
+     * @param maxBulkSize   the max size of each flush. The actual size is determined by the length of byte array.
+     * @param flushInterval if given maxBulkSize is not reached in this period, the flush would be trigger
+     *                      automatically. Unit is second.
+     * @param concurrency   the number of concurrency would run for the flush max.
+     * @param timeout       network timeout threshold in seconds.
+     * @return trace bulk write processor
+     */
+    public TraceBulkWriteProcessor buildTraceWriteProcessor(int maxBulkSize, int flushInterval, int concurrency, int timeout) {
+        checkState(this.traceServiceStub != null, "trace service is null");
+
+        return new TraceBulkWriteProcessor(this, maxBulkSize, flushInterval, concurrency, timeout, WRITE_HISTOGRAM, options);
+    }
+
+    /**
+     * Build a TraceWrite request.
+     *
+     * @param group     the group of the trace
+     * @param name      the name of the trace
+     * @param timestamp the timestamp of the trace
+     * @return the request to be built
+     */
+    public TraceWrite createTraceWrite(String group, String name, long timestamp) throws BanyanDBException {
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(group));
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(name));
+        return new TraceWrite(this.metadataCache.findTraceMetadata(group, name), timestamp);
+    }
+
+    /**
+     * Build a TraceWrite request without initial timestamp.
+     *
+     * @param group the group of the trace
+     * @param name  the name of the trace
+     * @return the request to be built
+     */
+    public TraceWrite createTraceWrite(String group, String name) throws BanyanDBException {
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(group));
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(name));
+        return new TraceWrite(this.metadataCache.findTraceMetadata(group, name));
+    }
+
     /**
      * Query streams according to given conditions
      *
@@ -456,6 +517,29 @@ public MeasureQueryResponse query(MeasureQuery measureQuery) throws BanyanDBExce
             throw new RuntimeException("No metadata found for the query");
    }
 
+    /**
+     * Query traces according to given conditions
+     *
+     * @param traceQuery condition for query
+     * @return trace query response.
+     */
+    public TraceQueryResponse query(TraceQuery traceQuery) throws BanyanDBException {
+        checkState(this.traceServiceStub != null, "trace service is null");
+
+        for (String group : traceQuery.groups) {
+            MetadataCache.EntityMetadata em = this.metadataCache.findTraceMetadata(group, traceQuery.name);
+            if (em != null) {
+                final BanyandbTrace.QueryResponse response = HandleExceptionsWith.callAndTranslateApiException(() ->
+                        this.traceServiceBlockingStub
+                                .withDeadlineAfter(this.getOptions().getDeadline(), TimeUnit.SECONDS)
+                                .query(traceQuery.build(em)));
+                return new TraceQueryResponse(response);
+            }
+
+        }
+        throw new RuntimeException("No metadata found for the query");
+    }
+
     /**
      * Define a new group and attach to the current client.
      *
@@ -968,6 +1052,69 @@ public DeleteResponse deleteProperty(String group, String name, String id) throw
         }
     }
 
+    /**
+     * Define a new trace
+     *
+     * @param trace the trace to be stored in the BanyanDB
+     * @throws BanyanDBException if the trace is invalid
+     */
+    public void define(Trace trace) throws BanyanDBException {
+        TraceMetadataRegistry registry = new TraceMetadataRegistry(checkNotNull(this.channel));
+        registry.create(trace);
+    }
+
+    /**
+     * Update the trace.
+     *
+     * @param trace the trace to be stored in the BanyanDB
+     * @throws BanyanDBException if the trace is invalid
+     */
+    public void update(Trace trace) throws BanyanDBException {
+        TraceMetadataRegistry registry = new TraceMetadataRegistry(checkNotNull(this.channel));
+        registry.update(trace);
+    }
+
+    /**
+     * Find the trace with given group and name
+     *
+     * @param group group of the metadata
+     * @param name  name of the metadata
+     * @return the trace found in BanyanDB. Otherwise, null is returned.
+     */
+    public Trace findTrace(String group, String name) throws BanyanDBException {
+        try {
+            return new TraceMetadataRegistry(checkNotNull(this.channel)).get(group, name);
+        } catch (BanyanDBException ex) {
+            if (ex.getStatus().equals(Status.Code.NOT_FOUND)) {
+                return null;
+            }
+            throw ex;
+        }
+    }
+
+    /**
+     * Find the traces with given group
+     *
+     * @param group group of the metadata
+     * @return the traces found in BanyanDB
+     */
+    public List<Trace> findTraces(String group) throws BanyanDBException {
+        TraceMetadataRegistry registry = new TraceMetadataRegistry(checkNotNull(this.channel));
+        return registry.list(group);
+    }
+
+    /**
+     * Delete the trace
+     *
+     * @param group group of the metadata
+     * @param name  name of the metadata
+     * @return if this trace has been deleted
+     */
+    public boolean deleteTrace(String group, String name) throws BanyanDBException {
+        TraceMetadataRegistry registry = new TraceMetadataRegistry(checkNotNull(this.channel));
+        return registry.delete(group, name);
+    }
+
     /**
      * Try to find the group defined
      *
@@ -1169,6 +1316,20 @@ public ResourceExist existProperty(String group, String name) throws BanyanDBExc
         return new PropertyMetadataRegistry(checkNotNull(this.channel)).exist(group, name);
     }
 
+    /**
+     * Check whether the trace definition is existed in the server
+     *
+     * @param group group of the metadata
+     * @param name  name of the metadata
+     * @return ResourceExist which indicates whether group and trace exist
+     */
+    public ResourceExist existTrace(String group, String name) throws BanyanDBException {
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(group));
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(name));
+
+        return new TraceMetadataRegistry(checkNotNull(this.channel)).exist(group, name);
+    }
+
     /**
      * Update the stream metadata cache from the server
      * @param group the group of the stream
@@ -1193,6 +1354,18 @@ public MetadataCache.EntityMetadata updateMeasureMetadataCacheFromSever(String g
         return this.metadataCache.updateMeasureFromSever(group, name);
     }
 
+    /**
+     * Update the trace metadata cache from the server
+     * @param group the group of the trace
+     * @param name the name of the trace
+     * @return the updated trace metadata, or null if the trace does not exist
+     */
+    public MetadataCache.EntityMetadata updateTraceMetadataCacheFromServer(String group, String name) throws BanyanDBException {
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(group));
+        Preconditions.checkArgument(!Strings.isNullOrEmpty(name));
+        return this.metadataCache.updateTraceFromServer(group, name);
+    }
+
     /**
      * Get the API version of the server
      *

src/main/java/org/apache/skywalking/banyandb/v1/client/TagAndValue.java

@@ -22,6 +22,7 @@
 
 import com.google.protobuf.ByteString;
 import com.google.protobuf.NullValue;
+import com.google.protobuf.Timestamp;
 import lombok.EqualsAndHashCode;
 import org.apache.skywalking.banyandb.model.v1.BanyandbModel;
 
@@ -77,6 +78,8 @@ public static TagAndValue<?> fromProtobuf(BanyandbModel.Tag tag) {
                 return new StringArrayTagPair(tag.getKey(), tag.getValue().getStrArray().getValueList());
             case BINARY_DATA:
                 return new BinaryTagPair(tag.getKey(), tag.getValue().getBinaryData());
+            case TIMESTAMP:
+                return new TimestampTagPair(tag.getKey(), tag.getValue().getTimestamp());
             case NULL:
                 return new NullTagPair(tag.getKey());
             default:
@@ -182,6 +185,32 @@ public static TagAndValue<ByteString> newBinaryTag(final String tagName, final b
         return new BinaryTagPair(tagName, ByteString.copyFrom(bytes));
     }
 
+    @EqualsAndHashCode(callSuper = true)
+    public static class TimestampTagPair extends TagAndValue<Long> {
+        TimestampTagPair(final String tagName, final long epochMilli) {
+            super(tagName, epochMilli);
+        }
+
+        TimestampTagPair(final String tagName, final Timestamp value) {
+            super(tagName, value.getSeconds() * 1000 + value.getNanos() / 1000000);
+        }
+
+        @Override
+        protected BanyandbModel.TagValue buildTypedTagValue() {
+            Timestamp timestamp = Timestamp.newBuilder()
+                    .setSeconds(value / 1000)
+                    .setNanos((int) ((value % 1000) * 1000000))
+                    .build();
+            return BanyandbModel.TagValue.newBuilder()
+                    .setTimestamp(timestamp)
+                    .build();
+        }
+    }
+
+    public static TagAndValue<Long> newTimestampTag(final String tagName, final long epochMilli) {
+        return new TimestampTagPair(tagName, epochMilli);
+    }
+
     @EqualsAndHashCode(callSuper = true)
     public static class NullTagPair extends TagAndValue<Void> {
         NullTagPair(final String tagName) {