[Notification Service] Introduces the countEvents interface

lib/airflow/airflow/www/views.py

@@ -2047,9 +2047,9 @@ def tree(self):
                             upstream_tasks[from_task_id] = {downstream_task}
 
         for key, value in event_tasks.items():
-            event_metas = self.notification_client.list_events(namespace=key[0], key=key[1], event_type=key[2])
-            for event_meta in event_metas:
-                sender = event_meta.sender
+            sender_event_counts = self.notification_client.count_events(namespace=key[0], key=key[1], event_type=key[2])
+            for sender_event_count in sender_event_counts[1]:
+                sender = sender_event_count.sender
                 if sender in dag.task_ids:
                     downstream_tasks = downstream_tasks.union(value)
                     if sender in upstream_tasks:
@@ -2225,7 +2225,7 @@ class GraphForm(DateTimeWithNumRunsWithDagRunsForm):
 
         dag_run = session.query(DagRun).filter(DagRun.dag_id == dag.dag_id, DagRun.execution_date == dttm).first()
 
-        task_instances:  Dict[str, Dict] = {}
+        task_instances: Dict[str, Dict] = {}
         task_logs: Dict[str, List[str]] = {}
         for ti in dag.get_task_instances(dttm, dttm):
             ti_dict = alchemy_to_dict(ti)
@@ -2262,13 +2262,14 @@ class GraphForm(DateTimeWithNumRunsWithDagRunsForm):
         send_events = {}
         event_senders = {}
         for key, value in events.items():
-            event_metas = self.notification_client.list_events(namespace=value[0], key=value[1], event_type=value[2])
-            event_senders[key] = event_metas
+            sender_event_counts = self.notification_client.count_events(namespace=value[0], key=value[1],
+                                                                        event_type=value[2])
+            event_senders[key] = sender_event_counts[1]
             sender_count = 0
-            for event_meta in event_metas:
-                if event_meta.sender in dag.task_ids:
-                    sender_count += 1
-            send_events[key] = (value[0], value[1], value[2], sender_count, len(event_metas) - sender_count)
+            for sender_event_count in sender_event_counts[1]:
+                if sender_event_count.sender in dag.task_ids:
+                    sender_count += sender_event_count.event_count
+            send_events[key] = (value[0], value[1], value[2], sender_count, sender_event_counts[0] - sender_count)
 
         node_edges = set()
         for e in edges:
@@ -2869,14 +2870,14 @@ def nodes(self, session=None):
         send_events = {}
         node_edges = set()
         for key, value in events.items():
-            event_metas = self.notification_client.list_events(namespace=value[0], key=value[1],
-                                                               event_type=value[2])
+            sender_event_counts = self.notification_client.count_events(namespace=value[0], key=value[1],
+                                                                        event_type=value[2])
             sender_count = 0
-            for event_meta in event_metas:
-                if event_meta.sender in dag.task_ids:
-                    sender_count += 1
-                    node_edges.add((event_meta.sender, key))
-            send_events[key] = (value[0], value[1], value[2], sender_count, len(event_metas) - sender_count)
+            for sender_event_count in sender_event_counts[1]:
+                if sender_event_count.sender in dag.task_ids:
+                    sender_count += sender_event_count.event_count
+                    node_edges.add((sender_event_count.sender, key))
+            send_events[key] = (value[0], value[1], value[2], sender_count, sender_event_counts[0] - sender_count)
 
         return json.dumps({'task_instances': task_instances, 'events': send_events,
                            'edges': [{'source_id': source_id, 'target_id': target_id} for source_id, target_id in
@@ -2897,17 +2898,13 @@ def event_senders(self):
         event_namespace = request.args.get('event_namespace')
         event_key = request.args.get('event_key')
         event_type = request.args.get('event_type')
-        events = self.notification_client.list_events(namespace=event_namespace, key=event_key, event_type=event_type)
-        logging.info('List events: {}'.format(str(events)))
+        sender_event_counts = self.notification_client.count_events(namespace=event_namespace, key=event_key,
+                                                                    event_type=event_type)
+        logging.info('Count events: {}'.format(str(sender_event_counts[1])))
         event_senders = {}
-        for event in events:
-            sender = event.sender
-            if sender in event_senders:
-                event_senders[sender] = event_senders[sender] + 1
-            else:
-                event_senders[sender] = 1
+        for sender_event_count in sender_event_counts[1]:
+            event_senders.update({sender_event_count.sender: sender_event_count.event_count})
         logging.info('List event senders: {}'.format(str(event_senders)))
-
         return json.dumps({'event_senders': event_senders},
                           cls=utils_json.AirflowJsonEncoder)
 

lib/notification_service/java/src/main/java/org/aiflow/notification/client/NotificationClient.java

@@ -24,9 +24,11 @@
 import io.grpc.ManagedChannelBuilder;
 import org.aiflow.notification.conf.Configuration;
 import org.aiflow.notification.entity.EventMeta;
+import org.aiflow.notification.entity.SenderEventCount;
 import org.aiflow.notification.proto.NotificationServiceGrpc;
 import org.aiflow.notification.proto.NotificationServiceOuterClass;
 import org.apache.commons.lang3.StringUtils;
+import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -224,6 +226,20 @@ private static List<EventMeta> parseEventsFromResponse(
         }
     }
 
+    private static ImmutablePair<Long, List<SenderEventCount>> parseEventCountFromResponse(
+            NotificationServiceOuterClass.CountEventsResponse response) throws Exception {
+        if (response.getReturnCode() == NotificationServiceOuterClass.ReturnStatus.SUCCESS) {
+            List<SenderEventCount> senderEventCounts = new ArrayList<>();
+            for (NotificationServiceOuterClass.SenderEventCountProto eventCountProto :
+                    response.getSenderEventCountsList()) {
+                senderEventCounts.add(SenderEventCount.buildSenderEventCount(eventCountProto));
+            }
+            return new ImmutablePair<>(response.getEventCount(), senderEventCounts);
+        } else {
+            throw new Exception(response.getReturnMsg());
+        }
+    }
+
     protected static NotificationServiceGrpc.NotificationServiceBlockingStub wrapBlockingStub(
             NotificationServiceGrpc.NotificationServiceBlockingStub stub,
             String target,
@@ -410,6 +426,37 @@ public List<EventMeta> listEvents(
                 0);
     }
 
+    /**
+     * Count specific `key` or `version` notifications in Notification Service.
+     *
+     * @param namespace Namespace of notification for listening.
+     * @param sender The sender of the event.
+     * @param keys Keys of notification for listening.
+     * @param version (Optional) Version of notification for listening.
+     * @param eventType (Optional) Type of event for listening.
+     * @param startTime (Optional) Type of event for listening.
+     * @return Count of Notification updated in Notification Service.
+     */
+    public ImmutablePair<Long, List<SenderEventCount>> countEvents(
+            String namespace,
+            List<String> keys,
+            long version,
+            String eventType,
+            long startTime,
+            String sender)
+            throws Exception {
+        NotificationServiceOuterClass.CountEventsRequest request =
+                NotificationServiceOuterClass.CountEventsRequest.newBuilder()
+                        .addAllKeys(keys)
+                        .setStartVersion(version)
+                        .setEventType(eventType)
+                        .setStartTime(startTime)
+                        .setNamespace(StringUtils.isEmpty(namespace) ? defaultNamespace : namespace)
+                        .setSender(sender)
+                        .build();
+        return parseEventCountFromResponse(notificationServiceStub.countEvents(request));
+    }
+
     /**
      * List all registered listener events in Notification Service.
      *

lib/notification_service/java/src/main/java/org/aiflow/notification/entity/SenderEventCount.java

@@ -0,0 +1,64 @@
+/*
+ * 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.aiflow.notification.entity;
+
+import org.aiflow.notification.proto.NotificationServiceOuterClass;
+
+public class SenderEventCount {
+
+    private String sender;
+    private long eventCount;
+
+    public SenderEventCount(String sender, long eventCount) {
+        this.sender = sender;
+        this.eventCount = eventCount;
+    }
+
+    public String getSender() {
+        return sender;
+    }
+
+    public void setSender(String sender) {
+        this.sender = sender;
+    }
+
+    public long getEventCount() {
+        return eventCount;
+    }
+
+    public void setEventCount(long eventCount) {
+        this.eventCount = eventCount;
+    }
+
+    @Override
+    public String toString() {
+        return "SenderEventCount{"
+                + "sender='"
+                + sender
+                + '\''
+                + ", eventCount="
+                + eventCount
+                + '}';
+    }
+
+    public static SenderEventCount buildSenderEventCount(
+            NotificationServiceOuterClass.SenderEventCountProto eventCountProto) {
+        return new SenderEventCount(eventCountProto.getSender(), eventCountProto.getEventCount());
+    }
+}

lib/notification_service/java/src/main/java/org/aiflow/notification/proto/NotificationServiceGrpc.java

@@ -113,6 +113,37 @@ org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsResponse>
     return getListEventsMethod;
   }
 
+  private static volatile io.grpc.MethodDescriptor<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest,
+      org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse> getCountEventsMethod;
+
+  @io.grpc.stub.annotations.RpcMethod(
+      fullMethodName = SERVICE_NAME + '/' + "countEvents",
+      requestType = org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest.class,
+      responseType = org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse.class,
+      methodType = io.grpc.MethodDescriptor.MethodType.UNARY)
+  public static io.grpc.MethodDescriptor<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest,
+      org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse> getCountEventsMethod() {
+    io.grpc.MethodDescriptor<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest, org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse> getCountEventsMethod;
+    if ((getCountEventsMethod = NotificationServiceGrpc.getCountEventsMethod) == null) {
+      synchronized (NotificationServiceGrpc.class) {
+        if ((getCountEventsMethod = NotificationServiceGrpc.getCountEventsMethod) == null) {
+          NotificationServiceGrpc.getCountEventsMethod = getCountEventsMethod =
+              io.grpc.MethodDescriptor.<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest, org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse>newBuilder()
+              .setType(io.grpc.MethodDescriptor.MethodType.UNARY)
+              .setFullMethodName(generateFullMethodName(SERVICE_NAME, "countEvents"))
+              .setSampledToLocalTracing(true)
+              .setRequestMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
+                  org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest.getDefaultInstance()))
+              .setResponseMarshaller(io.grpc.protobuf.ProtoUtils.marshaller(
+                  org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse.getDefaultInstance()))
+              .setSchemaDescriptor(new NotificationServiceMethodDescriptorSupplier("countEvents"))
+              .build();
+        }
+      }
+    }
+    return getCountEventsMethod;
+  }
+
   private static volatile io.grpc.MethodDescriptor<org.aiflow.notification.proto.NotificationServiceOuterClass.ListAllEventsRequest,
       org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsResponse> getListAllEventsMethod;
 
@@ -435,6 +466,16 @@ public void listEvents(org.aiflow.notification.proto.NotificationServiceOuterCla
       asyncUnimplementedUnaryCall(getListEventsMethod(), responseObserver);
     }
 
+    /**
+     * <pre>
+     * Count events.
+     * </pre>
+     */
+    public void countEvents(org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest request,
+        io.grpc.stub.StreamObserver<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse> responseObserver) {
+      asyncUnimplementedUnaryCall(getCountEventsMethod(), responseObserver);
+    }
+
     /**
      * <pre>
      * List all events
@@ -531,6 +572,13 @@ public void isClientExists(org.aiflow.notification.proto.NotificationServiceOute
                 org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsRequest,
                 org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsResponse>(
                   this, METHODID_LIST_EVENTS)))
+          .addMethod(
+            getCountEventsMethod(),
+            asyncUnaryCall(
+              new MethodHandlers<
+                org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest,
+                org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse>(
+                  this, METHODID_COUNT_EVENTS)))
           .addMethod(
             getListAllEventsMethod(),
             asyncUnaryCall(
@@ -633,6 +681,17 @@ public void listEvents(org.aiflow.notification.proto.NotificationServiceOuterCla
           getChannel().newCall(getListEventsMethod(), getCallOptions()), request, responseObserver);
     }
 
+    /**
+     * <pre>
+     * Count events.
+     * </pre>
+     */
+    public void countEvents(org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest request,
+        io.grpc.stub.StreamObserver<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse> responseObserver) {
+      asyncUnaryCall(
+          getChannel().newCall(getCountEventsMethod(), getCallOptions()), request, responseObserver);
+    }
+
     /**
      * <pre>
      * List all events
@@ -762,6 +821,16 @@ public org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsRes
           getChannel(), getListEventsMethod(), getCallOptions(), request);
     }
 
+    /**
+     * <pre>
+     * Count events.
+     * </pre>
+     */
+    public org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse countEvents(org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest request) {
+      return blockingUnaryCall(
+          getChannel(), getCountEventsMethod(), getCallOptions(), request);
+    }
+
     /**
      * <pre>
      * List all events
@@ -885,6 +954,17 @@ public com.google.common.util.concurrent.ListenableFuture<org.aiflow.notificatio
           getChannel().newCall(getListEventsMethod(), getCallOptions()), request);
     }
 
+    /**
+     * <pre>
+     * Count events.
+     * </pre>
+     */
+    public com.google.common.util.concurrent.ListenableFuture<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse> countEvents(
+        org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest request) {
+      return futureUnaryCall(
+          getChannel().newCall(getCountEventsMethod(), getCallOptions()), request);
+    }
+
     /**
      * <pre>
      * List all events
@@ -976,14 +1056,15 @@ public com.google.common.util.concurrent.ListenableFuture<org.aiflow.notificatio
 
   private static final int METHODID_SEND_EVENT = 0;
   private static final int METHODID_LIST_EVENTS = 1;
-  private static final int METHODID_LIST_ALL_EVENTS = 2;
-  private static final int METHODID_NOTIFY = 3;
-  private static final int METHODID_LIST_MEMBERS = 4;
-  private static final int METHODID_NOTIFY_NEW_MEMBER = 5;
-  private static final int METHODID_GET_LATEST_VERSION_BY_KEY = 6;
-  private static final int METHODID_REGISTER_CLIENT = 7;
-  private static final int METHODID_DELETE_CLIENT = 8;
-  private static final int METHODID_IS_CLIENT_EXISTS = 9;
+  private static final int METHODID_COUNT_EVENTS = 2;
+  private static final int METHODID_LIST_ALL_EVENTS = 3;
+  private static final int METHODID_NOTIFY = 4;
+  private static final int METHODID_LIST_MEMBERS = 5;
+  private static final int METHODID_NOTIFY_NEW_MEMBER = 6;
+  private static final int METHODID_GET_LATEST_VERSION_BY_KEY = 7;
+  private static final int METHODID_REGISTER_CLIENT = 8;
+  private static final int METHODID_DELETE_CLIENT = 9;
+  private static final int METHODID_IS_CLIENT_EXISTS = 10;
 
   private static final class MethodHandlers<Req, Resp> implements
       io.grpc.stub.ServerCalls.UnaryMethod<Req, Resp>,
@@ -1010,6 +1091,10 @@ public void invoke(Req request, io.grpc.stub.StreamObserver<Resp> responseObserv
           serviceImpl.listEvents((org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsRequest) request,
               (io.grpc.stub.StreamObserver<org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsResponse>) responseObserver);
           break;
+        case METHODID_COUNT_EVENTS:
+          serviceImpl.countEvents((org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsRequest) request,
+              (io.grpc.stub.StreamObserver<org.aiflow.notification.proto.NotificationServiceOuterClass.CountEventsResponse>) responseObserver);
+          break;
         case METHODID_LIST_ALL_EVENTS:
           serviceImpl.listAllEvents((org.aiflow.notification.proto.NotificationServiceOuterClass.ListAllEventsRequest) request,
               (io.grpc.stub.StreamObserver<org.aiflow.notification.proto.NotificationServiceOuterClass.ListEventsResponse>) responseObserver);
@@ -1105,6 +1190,7 @@ public static io.grpc.ServiceDescriptor getServiceDescriptor() {
               .setSchemaDescriptor(new NotificationServiceFileDescriptorSupplier())
               .addMethod(getSendEventMethod())
               .addMethod(getListEventsMethod())
+              .addMethod(getCountEventsMethod())
               .addMethod(getListAllEventsMethod())
               .addMethod(getNotifyMethod())
               .addMethod(getListMembersMethod())