[RemoteMirror] Fix demangle node corruption caused by excessively eager clearing of the NodeFactory.

mikeash · mikeash · commit fbc3f69ebf75 · 2023-03-29T16:59:21.000-04:00
We clear the NodeFactory to prevent unbounded buildup of allocated memory, but this is done too eagerly. In particular, normalizeReflectionName can end up clearing the factory while the calling code is still using nodes that were allocated from it.

To keep peak memory usage low while avoiding this problem, we introduce a checkpoint mechanism in NodeFactory. A checkpoint can be pushed and then subsequently popped. When a checkpoint is popped, only the nodes allocated since the checkpoint was pushed are invalidated and the memory reclaimed. This allows us to quickly clear short-lived nodes like those created in normalizeReflectionName, while preserving longer-lived nodes used in code calling it. Uses of clearNodeFactory are replaced with this checkpoint mechanism.

rdar://106547092
diff --git a/include/swift/Demangling/Demangler.h b/include/swift/Demangling/Demangler.h
@@ -53,7 +53,8 @@ class NodeFactory {
   /// The head of the single-linked slab list.
   Slab *CurrentSlab = nullptr;
 
-  /// The size of the previously allocated slab.
+  /// The size of the previously allocated slab. This may NOT be the size of
+  /// CurrentSlab, in the case where a checkpoint has been popped.
   ///
   /// The slab size can only grow, even clear() does not reset the slab size.
   /// This initial size is good enough to fit most de-manglings.
@@ -220,6 +221,26 @@ class NodeFactory {
     return {copiedString, str.size()};
   }
 
+  /// A checkpoint which captures the allocator's state at any given time. A
+  /// checkpoint can be popped to free all allocations made since it was made.
+  struct Checkpoint {
+    Slab *Slab;
+    char *CurPtr;
+    char *End;
+  };
+
+  /// Create a new checkpoint with the current state of the allocator.
+  Checkpoint pushCheckpoint() const;
+
+  /// Clear all allocations made since the given checkpoint. It is
+  /// undefined behavior to pop checkpoints in an order other than the
+  /// order in which they were pushed, or to pop a checkpoint when
+  /// clear() was called after creating it. The implementation attempts
+  /// to raise a fatal error in that case, but does not guarantee it. It
+  /// is allowed to pop outer checkpoints without popping inner ones, or
+  /// to call clear() without popping existing checkpoints.
+  void popCheckpoint(Checkpoint checkpoint);
+
   /// Creates a node of kind \p K.
   NodePointer createNode(Node::Kind K);
 
diff --git a/include/swift/RemoteInspection/TypeRefBuilder.h b/include/swift/RemoteInspection/TypeRefBuilder.h
@@ -456,7 +456,26 @@ class TypeRefBuilder {
 
   Demangle::NodeFactory &getNodeFactory() { return Dem; }
 
-  void clearNodeFactory() { Dem.clear(); }
+  NodeFactory::Checkpoint pushNodeFactoryCheckpoint() const {
+    return Dem.pushCheckpoint();
+  }
+
+  void popNodeFactoryCheckpoint(NodeFactory::Checkpoint checkpoint) {
+    Dem.popCheckpoint(checkpoint);
+  }
+
+  class ScopedNodeFactoryCheckpoint {
+    TypeRefBuilder *Builder;
+    NodeFactory::Checkpoint Checkpoint;
+
+  public:
+    ScopedNodeFactoryCheckpoint(TypeRefBuilder *Builder)
+        : Builder(Builder), Checkpoint(Builder->pushNodeFactoryCheckpoint()) {}
+
+    ~ScopedNodeFactoryCheckpoint() {
+      Builder->popNodeFactoryCheckpoint(Checkpoint);
+    }
+  };
 
   BuiltType decodeMangledType(Node *node, bool forRequirement = true);
 
@@ -1190,13 +1209,19 @@ class TypeRefBuilder {
       for (auto descriptor : sections.AssociatedType) {
         // Read out the relevant info from the associated type descriptor:
         // The type's name and which protocol conformance it corresponds to
-        auto typeRef = readTypeRef(descriptor, descriptor->ConformingTypeName);
-        auto typeName = nodeToString(demangleTypeRef(typeRef));
-        auto optionalMangledTypeName = normalizeReflectionName(typeRef);
-        auto protocolNode = demangleTypeRef(
-            readTypeRef(descriptor, descriptor->ProtocolTypeName));
-        auto protocolName = nodeToString(protocolNode);
-        clearNodeFactory();
+        llvm::Optional<std::string> optionalMangledTypeName;
+        std::string typeName;
+        std::string protocolName;
+        {
+          ScopedNodeFactoryCheckpoint checkpoint(this);
+          auto typeRef =
+              readTypeRef(descriptor, descriptor->ConformingTypeName);
+          typeName = nodeToString(demangleTypeRef(typeRef));
+          optionalMangledTypeName = normalizeReflectionName(typeRef);
+          auto protocolNode = demangleTypeRef(
+              readTypeRef(descriptor, descriptor->ProtocolTypeName));
+          protocolName = nodeToString(protocolNode);
+        }
         if (optionalMangledTypeName.has_value()) {
           auto mangledTypeName = optionalMangledTypeName.value();
           if (forMangledTypeName.has_value()) {
@@ -2011,10 +2036,10 @@ class TypeRefBuilder {
     std::unordered_map<std::string, std::string> typeNameToManglingMap;
     for (const auto &section : ReflectionInfos) {
       for (auto descriptor : section.Field) {
+        ScopedNodeFactoryCheckpoint checkpoint(this);
         auto TypeRef = readTypeRef(descriptor, descriptor->MangledTypeName);
         auto OptionalMangledTypeName = normalizeReflectionName(TypeRef);
         auto TypeName = nodeToString(demangleTypeRef(TypeRef));
-        clearNodeFactory();
         if (OptionalMangledTypeName.has_value()) {
           typeNameToManglingMap[TypeName] = OptionalMangledTypeName.value();
         }
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -505,13 +505,91 @@ void NodeFactory::clear() {
     // cycles.
     CurrentSlab->Previous = nullptr;
     CurPtr = (char *)(CurrentSlab + 1);
-    assert(End == CurPtr + SlabSize);
 #ifdef NODE_FACTORY_DEBUGGING
     allocatedMemory = 0;
 #endif
   }
 }
 
+NodeFactory::Checkpoint NodeFactory::pushCheckpoint() const {
+  return {CurrentSlab, CurPtr, End};
+}
+
+void NodeFactory::popCheckpoint(NodeFactory::Checkpoint checkpoint) {
+  if (checkpoint.Slab == CurrentSlab) {
+    if (checkpoint.CurPtr > CurPtr) {
+      fatal(0,
+            "Popping checkpoint {%p, %p, %p} that is after the current "
+            "pointer.\n",
+            checkpoint.Slab, checkpoint.CurPtr, checkpoint.End);
+    }
+    if (checkpoint.End != End) {
+      fatal(0,
+            "Popping checkpoint {%p, %p, %p} with End that does not match "
+            "current End %p.\n",
+            checkpoint.Slab, checkpoint.CurPtr, checkpoint.End, End);
+    }
+#ifndef NDEBUG
+    // Scribble the newly freed area.
+    memset(checkpoint.CurPtr, 0xaa, CurPtr - checkpoint.CurPtr);
+#endif
+    CurPtr = checkpoint.CurPtr;
+  } else {
+    // We may want to keep using the current slab rather than destroying
+    // it, since over time this allows us to converge on a steady state
+    // with no allocation activity (see the comment above in
+    // NodeFactory::clear). Keep using the current slab if the free space in the
+    // checkpoint's slab is less than 1/16th of the current slab's space. We
+    // won't repeatedly allocate and deallocate the current slab. The size
+    // doubles each time so we'll quickly pass the threshold.
+    Slab *savedSlab = nullptr;
+    if (CurrentSlab) {
+      size_t checkpointSlabFreeSpace = checkpoint.End - checkpoint.CurPtr;
+      size_t currentSlabSize = End - (char *)(CurrentSlab + 1);
+      if (currentSlabSize / 16 > checkpointSlabFreeSpace) {
+        savedSlab = CurrentSlab;
+        CurrentSlab = CurrentSlab->Previous;
+        // No need to save End, as it will still be in place later.
+      }
+    }
+
+    // Free all slabs (possibly except the one we saved) until we find the end
+    // or we find the checkpoint.
+    while (CurrentSlab && checkpoint.Slab != CurrentSlab) {
+      auto Slab = CurrentSlab;
+      CurrentSlab = CurrentSlab->Previous;
+      free(Slab);
+    }
+
+    // If we ran to the end and the checkpoint actually has a slab pointer, then
+    // the checkpoint is invalid.
+    if (!CurrentSlab && checkpoint.Slab) {
+      fatal(0,
+            "Popping checkpoint {%p, %p, %p} with slab that is not within "
+            "the allocator's slab chain.\n",
+            checkpoint.Slab, checkpoint.CurPtr, checkpoint.End);
+    }
+
+    if (savedSlab) {
+      // Reinstall the saved slab.
+      savedSlab->Previous = CurrentSlab;
+      CurrentSlab = savedSlab;
+      CurPtr = (char *)(CurrentSlab + 1);
+      // End is still valid from before.
+    } else {
+      // Set CurPtr and End to match the checkpoint's position.
+      CurPtr = checkpoint.CurPtr;
+      End = checkpoint.End;
+    }
+
+#ifndef NDEBUG
+    // Scribble the now freed end of the slab.
+    if (CurPtr)
+      memset(CurPtr, 0xaa, End - CurPtr);
+#endif
+  }
+}
+
 NodePointer NodeFactory::createNode(Node::Kind K) {
   return new (Allocate<Node>()) Node(K);
 }
diff --git a/stdlib/public/RemoteInspection/TypeRefBuilder.cpp b/stdlib/public/RemoteInspection/TypeRefBuilder.cpp
@@ -93,6 +93,7 @@ RemoteRef<char> TypeRefBuilder::readTypeRef(uint64_t remoteAddr) {
 /// Load and normalize a mangled name so it can be matched with string equality.
 llvm::Optional<std::string>
 TypeRefBuilder::normalizeReflectionName(RemoteRef<char> reflectionName) {
+  ScopedNodeFactoryCheckpoint checkpoint(this);
   // Remangle the reflection name to resolve symbolic references.
   if (auto node = demangleTypeRef(reflectionName,
                                   /*useOpaqueTypeSymbolicReferences*/ false)) {
@@ -104,7 +105,6 @@ TypeRefBuilder::normalizeReflectionName(RemoteRef<char> reflectionName) {
       return {};
     default:
       auto mangling = mangleNode(node);
-      clearNodeFactory();
       if (!mangling.isSuccess()) {
         return {};
       }
@@ -159,11 +159,11 @@ lookupTypeWitness(const std::string &MangledTypeName,
             0)
           continue;
 
+        ScopedNodeFactoryCheckpoint checkpoint(this);
         auto SubstitutedTypeName = readTypeRef(AssocTy,
                                                AssocTy->SubstitutedTypeName);
         auto Demangled = demangleTypeRef(SubstitutedTypeName);
         auto *TypeWitness = decodeMangledType(Demangled);
-        clearNodeFactory();
 
         AssociatedTypeCache.insert(std::make_pair(key, TypeWitness));
         return TypeWitness;
@@ -181,9 +181,9 @@ const TypeRef *TypeRefBuilder::lookupSuperclass(const TypeRef *TR) {
   if (!FD->hasSuperclass())
     return nullptr;
 
+  ScopedNodeFactoryCheckpoint checkpoint(this);
   auto Demangled = demangleTypeRef(readTypeRef(FD, FD->Superclass));
   auto Unsubstituted = decodeMangledType(Demangled);
-  clearNodeFactory();
   if (!Unsubstituted)
     return nullptr;
 
@@ -367,9 +367,9 @@ bool TypeRefBuilder::getFieldTypeRefs(
       continue;
     }
 
+    ScopedNodeFactoryCheckpoint checkpoint(this);
     auto Demangled = demangleTypeRef(readTypeRef(Field,Field->MangledTypeName));
     auto Unsubstituted = decodeMangledType(Demangled);
-    clearNodeFactory();
     if (!Unsubstituted)
       return false;
 
@@ -486,10 +486,10 @@ TypeRefBuilder::getClosureContextInfo(RemoteRef<CaptureDescriptor> CD) {
     auto CR = CD.getField(*i);
     
     if (CR->hasMangledTypeName()) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto MangledName = readTypeRef(CR, CR->MangledTypeName);
       auto DemangleTree = demangleTypeRef(MangledName);
       TR = decodeMangledType(DemangleTree);
-      clearNodeFactory();
     }
     Info.CaptureTypes.push_back(TR);
   }
@@ -499,10 +499,10 @@ TypeRefBuilder::getClosureContextInfo(RemoteRef<CaptureDescriptor> CD) {
     auto MSR = CD.getField(*i);
     
     if (MSR->hasMangledTypeName()) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto MangledName = readTypeRef(MSR, MSR->MangledTypeName);
       auto DemangleTree = demangleTypeRef(MangledName);
       TR = decodeMangledType(DemangleTree);
-      clearNodeFactory();
     }
 
     const MetadataSource *MS = nullptr;
@@ -526,11 +526,11 @@ TypeRefBuilder::getClosureContextInfo(RemoteRef<CaptureDescriptor> CD) {
 
 void TypeRefBuilder::dumpTypeRef(RemoteRef<char> MangledName,
                                  std::ostream &stream, bool printTypeName) {
+  ScopedNodeFactoryCheckpoint checkpoint(this);
   auto DemangleTree = demangleTypeRef(MangledName);
   auto TypeName = nodeToString(DemangleTree);
   stream << TypeName << "\n";
   auto Result = swift::Demangle::decodeMangledType(*this, DemangleTree);
-  clearNodeFactory();
   if (Result.isError()) {
     auto *Error = Result.getError();
     char *ErrorStr = Error->copyErrorString();
@@ -549,10 +549,14 @@ FieldTypeCollectionResult TypeRefBuilder::collectFieldTypes(
   FieldTypeCollectionResult result;
   for (const auto &sections : ReflectionInfos) {
     for (auto descriptor : sections.Field) {
-      auto typeRef = readTypeRef(descriptor, descriptor->MangledTypeName);
-      auto typeName = nodeToString(demangleTypeRef(typeRef));
-      auto optionalMangledTypeName = normalizeReflectionName(typeRef);
-      clearNodeFactory();
+      llvm::Optional<std::string> optionalMangledTypeName;
+      std::string typeName;
+      {
+        ScopedNodeFactoryCheckpoint checkpoint(this);
+        auto typeRef = readTypeRef(descriptor, descriptor->MangledTypeName);
+        typeName = nodeToString(demangleTypeRef(typeRef));
+        optionalMangledTypeName = normalizeReflectionName(typeRef);
+      }
       if (optionalMangledTypeName.has_value()) {
         auto mangledTypeName =
           optionalMangledTypeName.value();
@@ -618,10 +622,10 @@ void TypeRefBuilder::dumpFieldSection(std::ostream &stream) {
 void TypeRefBuilder::dumpBuiltinTypeSection(std::ostream &stream) {
   for (const auto &sections : ReflectionInfos) {
     for (auto descriptor : sections.Builtin) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto typeNode =
           demangleTypeRef(readTypeRef(descriptor, descriptor->TypeName));
       auto typeName = nodeToString(typeNode);
-      clearNodeFactory();
 
       stream << "\n- " << typeName << ":\n";
       stream << "Size: " << descriptor->Size << "\n";
@@ -670,10 +674,10 @@ void TypeRefBuilder::dumpCaptureSection(std::ostream &stream) {
 void TypeRefBuilder::dumpMultiPayloadEnumSection(std::ostream &stream) {
   for (const auto &sections : ReflectionInfos) {
     for (const auto descriptor : sections.MultiPayloadEnum) {
+      ScopedNodeFactoryCheckpoint checkpoint(this);
       auto typeNode =
           demangleTypeRef(readTypeRef(descriptor, descriptor->TypeName));
       auto typeName = nodeToString(typeNode);
-      clearNodeFactory();
 
       stream << "\n- " << typeName << ":\n";
       stream << "  Descriptor Size: " << descriptor->getSizeInBytes() << "\n";
