Use separate UsdStageCache's for stages that loadAll or loadNone.

lib/mayaUsd/fileio/jobs/readJob.cpp

@@ -101,16 +101,17 @@ UsdMaya_ReadJob::Read(std::vector<MDagPath>* addedDagPaths)
                                                                 varSelsVec);
 
     // Layer and Stage used to Read in the USD file
-    UsdStageCacheContext stageCacheContext(UsdMayaStageCache::Get());
     UsdStageRefPtr stage;
     if (mImportData.hasPopulationMask())
     {
+        // OpenMasked doesn't use the UsdStageCache, so don't create a UsdStageCacheContext
         stage = UsdStage::OpenMasked(rootLayer, sessionLayer,
                                      mImportData.stagePopulationMask(),
                                      mImportData.stageInitialLoadSet());
     }
     else
     {
+        UsdStageCacheContext stageCacheContext(UsdMayaStageCache::Get(mImportData.stageInitialLoadSet() == UsdStage::InitialLoadSet::LoadAll));
         stage = UsdStage::Open(rootLayer, sessionLayer,
                                mImportData.stageInitialLoadSet());
     }

lib/mayaUsd/nodes/proxyShapeBase.cpp

@@ -480,26 +480,32 @@ MayaUsdProxyShapeBase::computeInStageDataCached(MDataBlock& dataBlock)
             loadSet = UsdStage::InitialLoadSet::LoadNone;
         }
 
-        if (SdfLayerRefPtr rootLayer = SdfLayer::FindOrOpen(fileString)) {
-            UsdStageCacheContext ctx(UsdMayaStageCache::Get());
-            SdfLayerRefPtr sessionLayer = computeSessionLayer(dataBlock);
-            if (sessionLayer) {
-                usdStage = UsdStage::Open(rootLayer,
-                        sessionLayer,
-                        ArGetResolver().GetCurrentContext(),
-                        loadSet);
-            } else {
-                usdStage = UsdStage::Open(rootLayer,
-                        ArGetResolver().GetCurrentContext(),
-                        loadSet);
+        {
+            // When opening or creating stages we must have an active UsdStageCache.
+            // The stage cache is the only one who holds a strong reference to the
+            // UsdStage. See https://github.com/Autodesk/maya-usd/issues/528 for
+            // more information.
+            UsdStageCacheContext ctx(UsdMayaStageCache::Get(loadSet == UsdStage::InitialLoadSet::LoadAll));
+
+            if (SdfLayerRefPtr rootLayer = SdfLayer::FindOrOpen(fileString)) {
+                SdfLayerRefPtr sessionLayer = computeSessionLayer(dataBlock);
+                if (sessionLayer) {
+                    usdStage = UsdStage::Open(rootLayer,
+                            sessionLayer,
+                            ArGetResolver().GetCurrentContext(),
+                            loadSet);
+                } else {
+                    usdStage = UsdStage::Open(rootLayer,
+                            ArGetResolver().GetCurrentContext(),
+                            loadSet);
+                }
+
+                usdStage->SetEditTarget(usdStage->GetSessionLayer());
+            }
+            else {
+                // Create a new stage in memory with an anonymous root layer.
+                usdStage = UsdStage::CreateInMemory("", loadSet);
             }
-
-            usdStage->SetEditTarget(usdStage->GetSessionLayer());
-        }
-        else if (!usdStage) {
-            // Create a new stage in memory with an anonymous root layer.
-            UsdStageCacheContext ctx(UsdMayaStageCache::Get());
-            usdStage = UsdStage::CreateInMemory("", loadSet);
         }
 
         if (usdStage) {
@@ -755,6 +761,8 @@ MayaUsdProxyShapeBase::isStageValid() const
 MStatus
 MayaUsdProxyShapeBase::preEvaluation(const MDGContext& context, const MEvaluationNode& evaluationNode)
 {
+    // Any logic here should have an equivalent implementation in MayaUsdProxyShapeBase::setDependentsDirty().
+
     if (evaluationNode.dirtyPlugExists(excludePrimPathsAttr)) {
         _IncreaseExcludePrimPathsVersion();
     }
@@ -765,6 +773,7 @@ MayaUsdProxyShapeBase::preEvaluation(const MDGContext& context, const MEvaluatio
         evaluationNode.dirtyPlugExists(loadPayloadsAttr) ||
         evaluationNode.dirtyPlugExists(inStageDataAttr)) {
         _IncreaseUsdStageVersion();
+        MayaUsdProxyStageInvalidateNotice(*this).Send();
     }
 
     return MStatus::kSuccess;
@@ -786,6 +795,8 @@ MayaUsdProxyShapeBase::postEvaluation(const  MDGContext& context, const MEvaluat
 MStatus
 MayaUsdProxyShapeBase::setDependentsDirty(const MPlug& plug, MPlugArray& plugArray)
 {
+    // Any logic here should have an equivalent implementation in MayaUsdProxyShapeBase::preEvaluation() or postEvaluation().
+
     // If/when the MPxDrawOverride for the proxy shape specifies
     // isAlwaysDirty=false to improve performance, we must be sure to notify
     // the Maya renderer that the geometry is dirty and needs to be redrawn

lib/mayaUsd/nodes/stageData.h

@@ -104,6 +104,8 @@ class MayaUsdStageData : public MPxGeometryData
         // weak_ptr worked around this issue. Basically if we use a shared
         // pointer here, the only way to reload a scene, would be to restart
         // Maya."
+        //
+        // Logged as https://github.com/Autodesk/maya-usd/issues/528
 
         UsdStageWeakPtr stage;
         SdfPath primPath;

lib/mayaUsd/nodes/stageNode.cpp

@@ -123,7 +123,8 @@ UsdMayaStageNode::compute(const MPlug& plug, MDataBlock& dataBlock)
         UsdStageRefPtr usdStage;
 
         if (SdfLayerRefPtr rootLayer = SdfLayer::FindOrOpen(usdFile)) {
-            UsdStageCacheContext ctx(UsdMayaStageCache::Get());
+            const bool loadAll = true;
+            UsdStageCacheContext ctx(UsdMayaStageCache::Get(loadAll));
             usdStage = UsdStage::Open(rootLayer,
                                       ArGetResolver().GetCurrentContext());
 

lib/mayaUsd/python/wrapStageCache.cpp

@@ -33,7 +33,7 @@ void wrapStageCache()
     class_<UsdMayaStageCache>("StageCache")
 
         .def("Get", &UsdMayaStageCache::Get,
-             (args("forcePopulate") = true),
+             args("loadAll"),
              return_value_policy<reference_existing_object>())
         .staticmethod("Get")
         .def("Clear", &UsdMayaStageCache::Clear)

lib/mayaUsd/utils/stageCache.cpp

@@ -61,13 +61,13 @@ struct _OnSceneResetListener : public TfWeakBase {
 
 /* static */
 UsdStageCache&
-UsdMayaStageCache::Get(const bool forcePopulate)
+UsdMayaStageCache::Get(const bool loadAll)
 {
-    static UsdStageCache theCacheForcePopulate;
-    static UsdStageCache theCache;
+    static UsdStageCache theCacheLoadAll;  // used when UsdStage::Open() will be called with UsdStage::InitialLoadSet::LoadAll
+    static UsdStageCache theCache;         // used when UsdStage::Open() will be called with UsdStage::InitialLoadSet::LoadNode
     static _OnSceneResetListener onSceneResetListener;
 
-    return forcePopulate ? theCacheForcePopulate : theCache;
+    return loadAll ? theCacheLoadAll : theCache;
 }
 
 /* static */

lib/mayaUsd/utils/stageCache.h

@@ -31,10 +31,11 @@ class UsdMayaStageCache
 public:
 
     /// Return the singleton stage cache for use by all USD clients within Maya.
-    /// 2 stage caches are maintained; 1 for stages that have been
-    /// force-populated, and 1 for stages that have not been force-populated.
+    /// 2 stage caches are maintained; 1 for stages that have been opened with
+    /// UsdStage::InitialLoadSet::LoadAll, and 1 for stages that have been 
+    /// opened with UsdStage::InitialLoadSet::LoadNode.
     MAYAUSD_CORE_PUBLIC
-    static UsdStageCache& Get(const bool forcePopulate=true);
+    static UsdStageCache& Get(const bool loadAll);
 
     /// Clear the cache
     MAYAUSD_CORE_PUBLIC

plugin/pxr/maya/lib/usdMaya/referenceAssembly.cpp

@@ -732,7 +732,8 @@ UsdMayaReferenceAssembly::computeInStageDataCached(MDataBlock& dataBlock)
                 sessionLayer = unsharedSessionLayer;
             }
 
-            UsdStageCacheContext ctx(UsdMayaStageCache::Get());
+            const bool loadAll = true;
+            UsdStageCacheContext ctx(UsdMayaStageCache::Get(loadAll));
             usdStage = UsdStage::Open(rootLayer,
                                       sessionLayer,
                                       ArGetResolver().GetCurrentContext());
@@ -899,7 +900,8 @@ UsdMayaReferenceAssembly::computeOutStageData(MDataBlock& dataBlock)
             sessionLayer->TransferContent(oldLayer);
             sessionLayer->TransferContent(newLayer);
 
-            UsdStageCacheContext ctx(UsdMayaStageCache::Get());
+            const bool loadAll = true;
+            UsdStageCacheContext ctx(UsdMayaStageCache::Get(loadAll));
             usdStage = UsdStage::Open(usdPrim.GetStage()->GetRootLayer(),
                                       sessionLayer,
                                       ArGetResolver().GetCurrentContext());

plugin/pxr/maya/lib/usdMaya/translatorModelAssembly.cpp

@@ -375,7 +375,8 @@ UsdMayaTranslatorModelAssembly::Read(
         return false;
     }
 
-    UsdStageCacheContext stageCacheContext(UsdMayaStageCache::Get());
+    const bool loadAll = true;
+    UsdStageCacheContext stageCacheContext(UsdMayaStageCache::Get(loadAll));
     UsdStageRefPtr usdStage = UsdStage::Open(assetIdentifier);
     if (!usdStage) {
         TF_RUNTIME_ERROR("Cannot open USD file %s", assetIdentifier.c_str());

test/lib/testMayaUsdPythonImport.py

@@ -32,5 +32,5 @@ def testImportModule(self):
         # mayaUsd library also loaded its dependencies (i.e. the core USD
         # libraries). We test the type name as a string to ensure that we're
         # not causing USD libraries to be loaded any other way.
-        stageCache = mayaUsdLib.StageCache.Get()
+        stageCache = mayaUsdLib.StageCache.Get(True)
         self.assertEqual(type(stageCache).__name__, 'StageCache')