Merge pull request #1190 from Shopify/uk-disable-trackers-in-commands

Disable trackers in commands that don't need them

lib/tapioca.rb

@@ -11,16 +11,6 @@ module Tapioca
   class << self
     extend T::Sig
 
-    sig { params(trace_name: Symbol, block: T.proc.params(arg0: TracePoint).void).void }
-    def register_trace(trace_name, &block)
-      @traces << TracePoint.trace(trace_name, &block)
-    end
-
-    sig { void }
-    def disable_traces
-      @traces.each(&:disable)
-    end
-
     sig do
       type_parameters(:Result)
         .params(blk: T.proc.returns(T.type_parameter(:Result)))

lib/tapioca/cli.rb

@@ -269,8 +269,6 @@ def gem(*gems)
     option :payload, type: :boolean, desc: "Check shims against Sorbet's payload", default: true
     option :workers, aliases: ["-w"], type: :numeric, desc: "Number of parallel workers (default: auto)"
     def check_shims
-      Tapioca.disable_traces
-
       command = Commands::CheckShims.new(
         gem_rbi_dir: options[:gem_rbi_dir],
         dsl_rbi_dir: options[:dsl_rbi_dir],

lib/tapioca/commands.rb

@@ -4,6 +4,7 @@
 module Tapioca
   module Commands
     autoload :Command, "tapioca/commands/command"
+    autoload :CommandWithoutTracker, "tapioca/commands/command_without_tracker"
     autoload :Annotations, "tapioca/commands/annotations"
     autoload :CheckShims, "tapioca/commands/check_shims"
     autoload :Dsl, "tapioca/commands/dsl"

lib/tapioca/commands/annotations.rb

@@ -3,7 +3,7 @@
 
 module Tapioca
   module Commands
-    class Annotations < Command
+    class Annotations < CommandWithoutTracker
       extend T::Sig
 
       sig do

lib/tapioca/commands/check_shims.rb

@@ -3,7 +3,7 @@
 
 module Tapioca
   module Commands
-    class CheckShims < Command
+    class CheckShims < CommandWithoutTracker
       extend T::Sig
       include SorbetHelper
       include RBIFilesHelper

lib/tapioca/commands/command_without_tracker.rb

@@ -0,0 +1,18 @@
+# typed: strict
+# frozen_string_literal: true
+
+module Tapioca
+  module Commands
+    class CommandWithoutTracker < Command
+      extend T::Helpers
+
+      abstract!
+
+      sig { void }
+      def initialize
+        Tapioca::Runtime::Trackers.disable_all!
+        super
+      end
+    end
+  end
+end

lib/tapioca/commands/configure.rb

@@ -3,7 +3,7 @@
 
 module Tapioca
   module Commands
-    class Configure < Command
+    class Configure < CommandWithoutTracker
       sig do
         params(
           sorbet_config: String,

lib/tapioca/commands/dsl.rb

@@ -3,7 +3,7 @@
 
 module Tapioca
   module Commands
-    class Dsl < Command
+    class Dsl < CommandWithoutTracker
       include SorbetHelper
       include RBIFilesHelper
 

lib/tapioca/commands/require.rb

@@ -3,7 +3,7 @@
 
 module Tapioca
   module Commands
-    class Require < Command
+    class Require < CommandWithoutTracker
       sig do
         params(
           requires_path: String,

lib/tapioca/commands/todo.rb

@@ -3,7 +3,7 @@
 
 module Tapioca
   module Commands
-    class Todo < Command
+    class Todo < CommandWithoutTracker
       include SorbetHelper
 
       sig do

lib/tapioca/runtime/trackers.rb

@@ -1,6 +1,32 @@
-# typed: strict
+# typed: true
 # frozen_string_literal: true
 
+require "tapioca/runtime/trackers/tracker"
+
+module Tapioca
+  module Runtime
+    module Trackers
+      extend T::Sig
+
+      @trackers = T.let([], T::Array[Tracker])
+
+      class << self
+        extend T::Sig
+
+        sig { void }
+        def disable_all!
+          @trackers.each(&:disable!)
+        end
+
+        sig { params(tracker: Tracker).void }
+        def register_tracker(tracker)
+          @trackers << tracker
+        end
+      end
+    end
+  end
+end
+
 # The load order below is important:
 # ----------------------------------
 # We want the mixin tracker to be the first thing that is

lib/tapioca/runtime/trackers/autoload.rb

@@ -5,6 +5,7 @@ module Tapioca
   module Runtime
     module Trackers
       module Autoload
+        extend Tracker
         extend T::Sig
 
         NOOP_METHOD = ->(*_args, **_kwargs, &_block) {}
@@ -28,6 +29,8 @@ def eager_load_all!
 
           sig { params(constant_name: String).void }
           def register(constant_name)
+            return unless enabled?
+
             @constant_names_registered_for_autoload << constant_name
           end
 

lib/tapioca/runtime/trackers/constant_definition.rb

@@ -9,6 +9,7 @@ module Trackers
       # correspondence between classes/modules and files, as this information isn't
       # available in the ruby runtime without extra accounting.
       module ConstantDefinition
+        extend Tracker
         extend Reflection
         extend T::Sig
 
@@ -20,7 +21,7 @@ class ConstantLocation < T::Struct
         @class_files = {}.compare_by_identity
 
         # Immediately activated upon load. Observes class/module definition.
-        Tapioca.register_trace(:class) do |tp|
+        @class_tracepoint = TracePoint.trace(:class) do |tp|
           next if tp.self.singleton_class?
 
           key = tp.self
@@ -40,7 +41,7 @@ class ConstantLocation < T::Struct
           (@class_files[key] ||= Set.new) << loc
         end
 
-        Tapioca.register_trace(:c_return) do |tp|
+        @creturn_tracepoint = TracePoint.trace(:c_return) do |tp|
           next unless tp.method_id == :new
           next unless Module === tp.return_value
 
@@ -50,6 +51,12 @@ class ConstantLocation < T::Struct
         end
 
         class << self
+          def disable!
+            @class_tracepoint.disable
+            @creturn_tracepoint.disable
+            super
+          end
+
           def build_constant_location(tp, locations)
             file = resolve_loc(caller_locations)
             lineno = file == File.realpath(tp.path) ? tp.lineno : 0

lib/tapioca/runtime/trackers/mixin.rb

@@ -5,11 +5,11 @@ module Tapioca
   module Runtime
     module Trackers
       module Mixin
+        extend Tracker
         extend T::Sig
 
         @constants_to_mixin_locations = {}.compare_by_identity
         @mixins_to_constants = {}.compare_by_identity
-        @enabled = true
 
         class Type < T::Enum
           enums do
@@ -28,11 +28,7 @@ class << self
               .returns(T.type_parameter(:Result))
           end
           def with_disabled_registration(&block)
-            @enabled = false
-
-            block.call
-          ensure
-            @enabled = true
+            with_disabled_tracker(&block)
           end
 
           sig do
@@ -43,7 +39,7 @@ def with_disabled_registration(&block)
             ).void
           end
           def register(constant, mixin, mixin_type)
-            return unless @enabled
+            return unless enabled?
 
             location = Reflection.resolve_loc(caller_locations)
 

lib/tapioca/runtime/trackers/required_ancestor.rb

@@ -5,13 +5,16 @@ module Tapioca
   module Runtime
     module Trackers
       module RequiredAncestor
+        extend Tracker
         @required_ancestors_map = {}.compare_by_identity
 
         class << self
           extend T::Sig
 
           sig { params(requiring: T::Helpers, block: T.proc.void).void }
           def register(requiring, block)
+            return unless enabled?
+
             ancestors = @required_ancestors_map[requiring] ||= []
             ancestors << block
           end

lib/tapioca/runtime/trackers/tracker.rb

@@ -0,0 +1,45 @@
+# typed: true
+# frozen_string_literal: true
+
+module Tapioca
+  module Runtime
+    module Trackers
+      module Tracker
+        extend T::Sig
+        extend T::Helpers
+
+        abstract!
+
+        class << self
+          extend T::Sig
+
+          sig { params(base: T.all(Tracker, Module)).void }
+          def extended(base)
+            Trackers.register_tracker(base)
+            base.instance_exec do
+              @enabled = true
+            end
+          end
+        end
+
+        sig { void }
+        def disable!
+          @enabled = false
+        end
+
+        def enabled?
+          @enabled
+        end
+
+        def with_disabled_tracker(&block)
+          original_state = @enabled
+          @enabled = false
+
+          block.call
+        ensure
+          @enabled = original_state
+        end
+      end
+    end
+  end
+end

spec/tapioca/commands/annotations_spec.rb

@@ -7,6 +7,7 @@
 module Tapioca
   module Commands
     class AnnotationsSpec < SpecWithProject
+      include Tapioca::Helpers::Test::Isolation
       include WebMock::API
 
       DUMMY_REPO_URI_1 = "https://my-private-repo-1"