Use a compiler pass to register data collector conditionally

GromNaN · GromNaN · commit 950fa5294f9d · 2023-12-06T15:21:33.000+01:00
diff --git a/composer.json b/composer.json
@@ -28,6 +28,7 @@
         "symfony/framework-bundle": "^6.3.5 || ^7.0",
         "symfony/phpunit-bridge": "^6.3 || ^7.0",
         "symfony/yaml": "^6.3 || ^7.0",
+        "symfony/web-profiler-bundle": "^6.3 || ^7.0",
         "zenstruck/browser": "^1.6"
     },
     "scripts": {
diff --git a/config/services.php b/config/services.php
@@ -20,9 +20,9 @@
 
 namespace Symfony\Component\DependencyInjection\Loader\Configurator;
 
+use MongoDB\Bundle\Client;
 use MongoDB\Bundle\Command\DebugCommand;
 use MongoDB\Bundle\DataCollector\MongoDBDataCollector;
-use MongoDB\Bundle\TraceableClient;
 
 return static function (ContainerConfigurator $container): void {
     // default configuration for services in *this* file
@@ -39,14 +39,13 @@
         ->tag('console.command');
 
     $services
-        ->set('mongodb.prototype.client', TraceableClient::class)
+        ->set('mongodb.prototype.client', Client::class)
         ->arg('$uri', abstract_arg('Should be defined by pass'))
         ->arg('$uriOptions', abstract_arg('Should be defined by pass'))
-        ->arg('$driverOptions', abstract_arg('Should be defined by pass'))
-        ->tag('mongodb.client');
+        ->arg('$driverOptions', abstract_arg('Should be defined by pass'));
 
     $services
-        ->set('data_collector.mongodb', MongoDBDataCollector::class)
+        ->set('mongodb.data_collector', MongoDBDataCollector::class)
         ->tag('data_collector', [
             'template' => '@MongoDB/Collector/mongodb.html.twig',
             'id' => 'mongodb',
diff --git a/src/Client.php b/src/Client.php
@@ -0,0 +1,35 @@
+<?php
+
+declare(strict_types=1);
+
+/**
+ * Copyright 2023-present MongoDB, Inc.
+ *
+ * Licensed 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
+ *
+ *   https://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.
+ */
+
+namespace MongoDB\Bundle;
+
+use MongoDB\Client as MongoDBClient;
+use MongoDB\Driver\Monitoring\Subscriber;
+
+final class Client extends MongoDBClient
+{
+    /**
+     * @internal
+     */
+    public function addSubscriber(Subscriber $subscriber): void
+    {
+        $this->getManager()->addSubscriber($subscriber);
+    }
+}
diff --git a/src/DataCollector/MongoDBDataCollector.php b/src/DataCollector/MongoDBDataCollector.php
@@ -20,9 +20,7 @@
 
 namespace MongoDB\Bundle\DataCollector;
 
-use MongoDB\BSON\Document;
-use MongoDB\Bundle\TraceableClient;
-use MongoDB\Client;
+use MongoDB\Bundle\Client;
 use MongoDB\Driver\Monitoring\CommandFailedEvent;
 use MongoDB\Driver\Monitoring\CommandStartedEvent;
 use MongoDB\Driver\Monitoring\CommandSucceededEvent;
@@ -33,67 +31,61 @@
 final class MongoDBDataCollector extends DataCollector
 {
     /**
-     * @var list<Client>
+     * @var list<array{client:Client, subscriber:DriverEventSubscriber}>
      */
     private array $clients = [];
 
-    public function addClient(string $name, Client $client): void
+    public function addClient(string $name, Client $client, DriverEventSubscriber $subscriber): void
     {
-        $this->clients[$name] = $client;
+        $this->clients[$name] = [
+            'client' => $client,
+            'subscriber' => $subscriber,
+        ];
     }
 
     public function collect(Request $request, Response $response, ?\Throwable $exception = null): void
     {
-        foreach ($this->clients as $name => $client) {
+        foreach ($this->clients as $name => ['client' => $client, 'subscriber' => $subscriber]) {
             $totalTime = 0;
             $requestCount = 0;
             $errorCount = 0;
             $requests = [];
 
-            if ($client instanceof TraceableClient) {
-                foreach ($client->getEvents() as $event) {
-                    $requestId = $event->getRequestId();
+            foreach ($subscriber->getEvents() as $event) {
+                $requestId = $event->getRequestId();
 
-                    $eventData = [
-                        'class' => $event::class,
-                        'commandName' => $event->getCommandName(),
-                        'server' => $event->getServer()->getInfo(),
+                if ($event instanceof CommandStartedEvent) {
+                    $command = (array) $event->getCommand();
+                    unset($command['lsid'], $command['$clusterTime']);
+
+                    $requests[$requestId] = [
                         'client' => $name,
+                        'startedAt' => hrtime(true),
+                        'commandName' => $event->getCommandName(),
+                        'command' => $command,
+                        // 'server' => $event->getServer()->getInfo(),
+                        'operationId' => $event->getOperationId(),
+                        'database' => $event->getDatabaseName(),
+                        'serviceId' => $event->getServiceId(),
                     ];
-
-                    if ($event instanceof CommandStartedEvent) {
-                        $command = (array) $event->getCommand();
-                        unset($command['lsid'], $command['$clusterTime']);
-
-                        $requests[$requestId] = [
-                            'client' => $name,
-                            'startedAt' => hrtime(true),
-                            'commandName' => $event->getCommandName(),
-                            'command' => $command,
-                            // 'server' => $event->getServer()->getInfo(),
-                            'operationId' => $event->getOperationId(),
-                            'database' => $event->getDatabaseName(),
-                            'serviceId' => $event->getServiceId(),
-                        ];
-                        ++$requestCount;
-                    } elseif ($event instanceof CommandSucceededEvent) {
-                        $requests[$requestId] += [
-                            // 'reply' => Document::fromPHP($event->getReply()),
-                            'duration' => $event->getDurationMicros(),
-                            'endedAt' => hrtime(true),
-                            'success' => true,
-                        ];
-                        $totalTime += $event->getDurationMicros();
-                    } elseif ($event instanceof CommandFailedEvent) {
-                        $requests[$requestId] += [
-                            // 'reply' => Document::fromPHP($event->getReply()),
-                            'duration' => $event->getDurationMicros(),
-                            'error' => $event->getError(),
-                            'success' => false,
-                        ];
-                        $totalTime += $event->getDurationMicros();
-                        ++$errorCount;
-                    }
+                    ++$requestCount;
+                } elseif ($event instanceof CommandSucceededEvent) {
+                    $requests[$requestId] += [
+                        // 'reply' => Document::fromPHP($event->getReply()),
+                        'duration' => $event->getDurationMicros(),
+                        'endedAt' => hrtime(true),
+                        'success' => true,
+                    ];
+                    $totalTime += $event->getDurationMicros();
+                } elseif ($event instanceof CommandFailedEvent) {
+                    $requests[$requestId] += [
+                        // 'reply' => Document::fromPHP($event->getReply()),
+                        'duration' => $event->getDurationMicros(),
+                        'error' => $event->getError(),
+                        'success' => false,
+                    ];
+                    $totalTime += $event->getDurationMicros();
+                    ++$errorCount;
                 }
             }
 
diff --git a/src/DependencyInjection/Compiler/DataCollectorPass.php b/src/DependencyInjection/Compiler/DataCollectorPass.php
@@ -0,0 +1,52 @@
+<?php
+
+declare(strict_types=1);
+
+/**
+ * Copyright 2023-present MongoDB, Inc.
+ *
+ * Licensed 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
+ *
+ *   https://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.
+ */
+
+namespace MongoDB\Bundle\DependencyInjection\Compiler;
+
+use MongoDB\Bundle\DataCollector\DriverEventSubscriber;
+use Symfony\Component\DependencyInjection\Compiler\CompilerPassInterface;
+use Symfony\Component\DependencyInjection\ContainerBuilder;
+use Symfony\Component\DependencyInjection\Definition;
+use Symfony\Component\DependencyInjection\Reference;
+
+/** @internal */
+final class DataCollectorPass implements CompilerPassInterface
+{
+    public function process(ContainerBuilder $container): void
+    {
+        if (!$container->has('profiler')) {
+            return;
+        }
+
+        $dataCollector = $container->getDefinition('mongodb.data_collector');
+
+        // Add a subscriber to each client to collect driver events, and register the client to the data collector.
+        foreach ($container->findTaggedServiceIds('mongodb.client', true) as $clientId => $attributes) {
+            $subscriberId = sprintf('%s.subscriber', $clientId);
+            $container->setDefinition($subscriberId, new Definition(DriverEventSubscriber::class));
+            $container->getDefinition($clientId)->addMethodCall('addSubscriber', [new Reference($subscriberId)]);
+            $dataCollector->addMethodCall('addClient', [
+                $attributes[0]['name'] ?? $clientId,
+                new Reference($clientId),
+                new Reference($subscriberId),
+            ]);
+        }
+    }
+}
diff --git a/src/DependencyInjection/MongoDBExtension.php b/src/DependencyInjection/MongoDBExtension.php
@@ -27,7 +27,6 @@
 use Symfony\Component\DependencyInjection\Definition;
 use Symfony\Component\DependencyInjection\Extension\Extension;
 use Symfony\Component\DependencyInjection\Loader\PhpFileLoader;
-use Symfony\Component\DependencyInjection\Reference;
 
 use function dirname;
 use function sprintf;
@@ -73,11 +72,9 @@ private function createClients(string $defaultClient, array $clients, ContainerB
             $clientDefinition->setArgument('$uri', $configuration['uri']);
             $clientDefinition->setArgument('$uriOptions', $configuration['uriOptions'] ?? []);
             $clientDefinition->setArgument('$driverOptions', $configuration['driverOptions'] ?? []);
-
+            $clientDefinition->addTag('mongodb.client', ['name' => $client]);
             $container->setDefinition($serviceId, $clientDefinition);
 
-            $dataCollector->addMethodCall('addClient', [$client, new Reference($serviceId)]);
-
             if (isset($configuration['default_database'])) {
                 $container->setParameter(sprintf('%s.default_database', $serviceId), $configuration['default_database']);
             }
diff --git a/src/MongoDBBundle.php b/src/MongoDBBundle.php
@@ -21,6 +21,7 @@
 namespace MongoDB\Bundle;
 
 use MongoDB\Bundle\DependencyInjection\MongoDBExtension;
+use Symfony\Component\DependencyInjection\ContainerBuilder;
 use Symfony\Component\DependencyInjection\Extension\ExtensionInterface;
 use Symfony\Component\HttpKernel\Bundle\AbstractBundle;
 
@@ -30,4 +31,9 @@ public function getContainerExtension(): ?ExtensionInterface
     {
         return new MongoDBExtension();
     }
+
+    public function build(ContainerBuilder $container): void
+    {
+        $container->addCompilerPass(new DependencyInjection\Compiler\DataCollectorPass());
+    }
 }
diff --git a/src/TraceableClient.php b/src/TraceableClient.php
diff --git a/tests/Unit/DependencyInjection/MongoDBExtensionTest.php b/tests/Unit/DependencyInjection/MongoDBExtensionTest.php
@@ -59,7 +59,7 @@ public function testLoadWithSingleClient(): void
 
         // Check service definition
         $definition = $container->getDefinition('mongodb.client.default');
-        $this->assertSame(Client::class, $definition->getClass());
+        $this->assertTrue(is_a($definition->getClass(), Client::class, true), sprintf('Expected "%s" to be a "%s"', $definition->getClass(), Client::class));
         $this->assertSame('mongodb://localhost:27017', $definition->getArgument('$uri'));
 
         // Check alias definition
@@ -93,12 +93,12 @@ public function testLoadWithMultipleClients(): void
 
         // Check service definitions
         $definition = $container->getDefinition('mongodb.client.default');
-        $this->assertSame(Client::class, $definition->getClass());
+        $this->assertTrue(is_a($definition->getClass(), Client::class, true), sprintf('Expected "%s" to be a "%s"', $definition->getClass(), Client::class));
         $this->assertSame('mongodb://localhost:27017', $definition->getArgument('$uri'));
         $this->assertSame(['readPreference' => 'primary'], $definition->getArgument('$uriOptions'));
 
         $definition = $container->getDefinition('mongodb.client.secondary');
-        $this->assertSame(Client::class, $definition->getClass());
+        $this->assertTrue(is_a($definition->getClass(), Client::class, true), sprintf('Expected "%s" to be a "%s"', $definition->getClass(), Client::class));
         $this->assertSame('mongodb://localhost:27018', $definition->getArgument('$uri'));
         $this->assertEquals(['serverApi' => new ServerApi((string) ServerApi::V1)], $definition->getArgument('$driverOptions'));
     }

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@`
`21`	`21`	`namespace MongoDB\Bundle;`
`22`	`22`
`23`	`23`	`use MongoDB\Bundle\DependencyInjection\MongoDBExtension;`
	`24`	`+use Symfony\Component\DependencyInjection\ContainerBuilder;`
`24`	`25`	`use Symfony\Component\DependencyInjection\Extension\ExtensionInterface;`
`25`	`26`	`use Symfony\Component\HttpKernel\Bundle\AbstractBundle;`
`26`	`27`
`@@ -30,4 +31,9 @@ public function getContainerExtension(): ?ExtensionInterface`
`30`	`31`	`{`
`31`	`32`	`return new MongoDBExtension();`
`32`	`33`	`}`
	`34`	`+`
	`35`	`+ public function build(ContainerBuilder $container): void`
	`36`	`+ {`
	`37`	`+ $container->addCompilerPass(new DependencyInjection\Compiler\DataCollectorPass());`
	`38`	`+ }`
`33`	`39`	`}`