Allow for multiple copies of the same function in different modules.

Previous iterations of the cache assumed a one-one-one function name to
module mapping, however it's possible to end up with many modules which
contain the same function (e.g due to a race condition when using multiple
processes, but other scenarios could exist in the future). This commit
separates the func_name (the id for the function) with the link_name
(the unique id for a given function in a given module). This distinction
is not exposed outside the cache - when asked for a target the cache
chooses which version to return (currently just the first one it sees).

Author: William Grant

typed_python/compiler/compiler_cache.py

@@ -49,77 +49,88 @@ class CompilerCache:
     when we first boot up, which could be slow. We could improve this substantially
     by making it possible to determine if a given function is in the cache by organizing
     the manifests by, say, function name.
+
+    Due to the potential for race conditions, we must distinguish between the following:
+        func_name -  The identifier for the function, based on its identity hash.
+        link_name -  The identifier for the specific realization of that function, which lives in a specific
+            cache module.
     """
     def __init__(self, cacheDir):
         self.cacheDir = cacheDir
 
         ensureDirExists(cacheDir)
 
         self.loadedBinarySharedObjects = Dict(str, LoadedBinarySharedObject)()
-        self.nameToModuleHash = Dict(str, str)()
-
+        self.link_name_to_module_hash = Dict(str, str)()
         self.moduleManifestsLoaded = set()
-
+        # link_names with an associated module in loadedBinarySharedObjects
+        self.targetsLoaded: Dict[str, TypedCallTarget] = {}
+        # the set of link_names for functions with linked and validated globals (i.e. ready to be run).
+        self.targetsValidated = set()
+        # link_name -> link_name
+        self.function_dependency_graph = DirectedGraph()
+        # dict from link_name to list of global names (should be llvm keys in serialisedGlobalDefinitions)
+        self.global_dependencies = Dict(str, ListOf(str))()
+        self.func_name_to_link_names = Dict(str, ListOf(str))()
         for moduleHash in os.listdir(self.cacheDir):
             if len(moduleHash) == 40:
                 self.loadNameManifestFromStoredModuleByHash(moduleHash)
 
-        # the set of functions with an associated module in loadedBinarySharedObjects
-        self.targetsLoaded: Dict[str, TypedCallTarget] = {}
+    def hasSymbol(self, func_name: str) -> bool:
+        """Returns true if there are any versions of `func_name` in the cache.
 
-        # the set of functions with linked and validated globals (i.e. ready to be run).
-        self.targetsValidated = set()
+        There may be multiple copies in different modules with different link_names.
+        """
+        return any(link_name in self.link_name_to_module_hash for link_name in self.func_name_to_link_names.get(func_name, []))
 
-        self.function_dependency_graph = DirectedGraph()
-        # dict from function linkname to list of global names (should be llvm keys in serialisedGlobalDefinitions)
-        self.global_dependencies = Dict(str, ListOf(str))()
+    def getTarget(self, func_name: str) -> TypedCallTarget:
+        if not self.hasSymbol(func_name):
+            raise ValueError(f'symbol not found for func_name {func_name}')
+        link_name = self._select_link_name(func_name)
+        self.loadForSymbol(link_name)
+        return self.targetsLoaded[link_name]
 
-    def hasSymbol(self, linkName: str) -> bool:
-        """NB this will return True even if the linkName is ultimately unretrievable."""
-        return linkName in self.nameToModuleHash
+    def _generate_link_name(self, func_name: str, module_hash: str) -> str:
+        return func_name + "." + module_hash
 
-    def getTarget(self, linkName: str) -> TypedCallTarget:
-        if not self.hasSymbol(linkName):
-            raise ValueError(f'symbol not found for linkName {linkName}')
-        self.loadForSymbol(linkName)
-        return self.targetsLoaded[linkName]
+    def _select_link_name(self, func_name) -> str:
+        """choose a link name for a given func name.
 
-    def dependencies(self, linkName: str) -> Optional[List[str]]:
-        """Returns all the function names that `linkName` depends on"""
-        return list(self.function_dependency_graph.outgoing(linkName))
+        Currently we just choose the first available option.
+        Throws a KeyError if func_name isn't in the cache.
+        """
+        link_name_candidates = self.func_name_to_link_names[func_name]
+        return link_name_candidates[0]
+
+    def dependencies(self, link_name: str) -> Optional[List[str]]:
+        """Returns all the function names that `link_name` depends on"""
+        return list(self.function_dependency_graph.outgoing(link_name))
 
     def loadForSymbol(self, linkName: str) -> None:
-        """Loads the whole module, and any submodules, into LoadedBinarySharedObjects"""
-        moduleHash = self.nameToModuleHash[linkName]
+        """Loads the whole module, and any dependant modules, into LoadedBinarySharedObjects"""
+        moduleHash = self.link_name_to_module_hash[linkName]
 
         self.loadModuleByHash(moduleHash)
 
         if linkName not in self.targetsValidated:
-            dependantFuncs = self.dependencies(linkName) + [linkName]
-            globalsToLink = {}  # dict from modulehash to list of globals.
-            for funcName in dependantFuncs:
-                if funcName not in self.targetsValidated:
-                    funcModuleHash = self.nameToModuleHash[funcName]
-                    # append to the list of globals to link for a given module.  TODO: optimise this, don't double-link.
-                    globalsToLink[funcModuleHash] = globalsToLink.get(funcModuleHash, []) + self.global_dependencies.get(funcName, [])
-
-            for moduleHash, globs in globalsToLink.items():  # this works because loadModuleByHash loads submodules too.
-                if globs:
-                    definitionsToLink = {x: self.loadedBinarySharedObjects[moduleHash].serializedGlobalVariableDefinitions[x]
-                                         for x in globs
-                                         }
-                    self.loadedBinarySharedObjects[moduleHash].linkGlobalVariables(definitionsToLink)
-                    if not self.loadedBinarySharedObjects[moduleHash].validateGlobalVariables(definitionsToLink):
-                        raise RuntimeError('failed to validate globals when loading:', linkName)
-
-            self.targetsValidated.update(dependantFuncs)
+            self.targetsValidated.add(linkName)
+            for dependant_func in self.dependencies(linkName):
+                self.loadForSymbol(dependant_func)
+
+            globalsToLink = self.global_dependencies.get(linkName, [])
+            if globalsToLink:
+                definitionsToLink = {x: self.loadedBinarySharedObjects[moduleHash].serializedGlobalVariableDefinitions[x]
+                                     for x in globalsToLink
+                                     }
+                self.loadedBinarySharedObjects[moduleHash].linkGlobalVariables(definitionsToLink)
+                if not self.loadedBinarySharedObjects[moduleHash].validateGlobalVariables(definitionsToLink):
+                    raise RuntimeError('failed to validate globals when loading:', linkName)
 
     def loadModuleByHash(self, moduleHash: str) -> None:
         """Load a module by name.
 
-        As we load, place all the newly imported typed call targets into
-        'nameToTypedCallTarget' so that the rest of the system knows what functions
-        have been uncovered.
+        Add the module contents to targetsLoaded, generate a LoadedBinarySharedObject,
+        and update the function and global dependency graphs.
         """
         if moduleHash in self.loadedBinarySharedObjects:
             return
@@ -128,6 +139,7 @@ def loadModuleByHash(self, moduleHash: str) -> None:
 
         # TODO (Will) - store these names as module consts, use one .dat only
         with open(os.path.join(targetDir, "type_manifest.dat"), "rb") as f:
+            # func_name -> typedcalltarget
             callTargets = SerializationContext().deserialize(f.read())
 
         with open(os.path.join(targetDir, "globals_manifest.dat"), "rb") as f:
@@ -156,45 +168,68 @@ def loadModuleByHash(self, moduleHash: str) -> None:
             serializedGlobalVarDefs,
             functionNameToNativeType,
             globalDependencies
-
         ).loadFromPath(modulePath)
 
         self.loadedBinarySharedObjects[moduleHash] = loaded
 
-        self.targetsLoaded.update(callTargets)
+        for func_name, callTarget in callTargets.items():
+            link_name = self._generate_link_name(func_name, moduleHash)
+            assert link_name not in self.targetsLoaded
+            self.targetsLoaded[link_name] = callTarget
 
-        assert not any(key in self.global_dependencies for key in globalDependencies)  # should only happen if there's a hash collision.
-        self.global_dependencies.update(globalDependencies)
+        link_name_global_dependencies = {self._generate_link_name(x, moduleHash): y for x, y in globalDependencies.items()}
 
+        assert not any(key in self.global_dependencies for key in link_name_global_dependencies)
+
+        self.global_dependencies.update(link_name_global_dependencies)
         # update the cache's dependency graph with our new edges.
         for function_name, dependant_function_name in dependency_edgelist:
             self.function_dependency_graph.addEdge(source=function_name, dest=dependant_function_name)
 
     def addModule(self, binarySharedObject, nameToTypedCallTarget, linkDependencies, dependencyEdgelist):
         """Add new code to the compiler cache.
 
+
         Args:
             binarySharedObject: a BinarySharedObject containing the actual assembler
                 we've compiled.
-            nameToTypedCallTarget: a dict from linkname to TypedCallTarget telling us
+            nameToTypedCallTarget: a dict from func_name to TypedCallTarget telling us
                 the formal python types for all the objects.
-            linkDependencies: a set of linknames we depend on directly.
+            linkDependencies: a set of func_names we depend on directly. (this becomes submodules)
             dependencyEdgelist (list): a list of source, dest pairs giving the set of dependency graph for the
                 module.
+
+        TODO (Will): the notion of submodules/linkDependencies can be refactored out.
         """
-        dependentHashes = set()
 
+        hashToUse = SerializationContext().sha_hash(str(uuid.uuid4())).hexdigest
+
+        # the linkDependencies and dependencyEdgelist are in terms of func_name.
+        dependentHashes = set()
         for name in linkDependencies:
-            dependentHashes.add(self.nameToModuleHash[name])
+            link_name = self._select_link_name(name)
+            dependentHashes.add(self.link_name_to_module_hash[link_name])
+
+        link_name_dependency_edgelist = []
+        for source, dest in dependencyEdgelist:
+            assert source in binarySharedObject.definedSymbols
+            source_link_name = self._generate_link_name(source, hashToUse)
+            if dest in binarySharedObject.definedSymbols:
+                dest_link_name = self._generate_link_name(dest, hashToUse)
+            else:
+                dest_link_name = self._select_link_name(dest)
+            link_name_dependency_edgelist.append([source_link_name, dest_link_name])
 
-        path, hashToUse = self.writeModuleToDisk(binarySharedObject, nameToTypedCallTarget, dependentHashes, dependencyEdgelist)
+        path = self.writeModuleToDisk(binarySharedObject, hashToUse, nameToTypedCallTarget, dependentHashes, link_name_dependency_edgelist)
 
         self.loadedBinarySharedObjects[hashToUse] = (
             binarySharedObject.loadFromPath(os.path.join(path, "module.so"))
         )
 
-        for n in binarySharedObject.definedSymbols:
-            self.nameToModuleHash[n] = hashToUse
+        for func_name in binarySharedObject.definedSymbols:
+            link_name = self._generate_link_name(func_name, hashToUse)
+            self.link_name_to_module_hash[link_name] = hashToUse
+            self.func_name_to_link_names.setdefault(func_name, []).append(link_name)
 
         # link & validate all globals for the new module
         self.loadedBinarySharedObjects[hashToUse].linkGlobalVariables()
@@ -208,20 +243,18 @@ def loadNameManifestFromStoredModuleByHash(self, moduleHash) -> None:
 
         targetDir = os.path.join(self.cacheDir, moduleHash)
 
-        with open(os.path.join(targetDir, "submodules.dat"), "rb") as f:
-            submodules = SerializationContext().deserialize(f.read(), ListOf(str))
-
-        for subHash in submodules:
-            self.loadNameManifestFromStoredModuleByHash(subHash)
-
+        # TODO (Will) the name_manifest module_hash is the same throughout so this doesn't need to be a dict.
         with open(os.path.join(targetDir, "name_manifest.dat"), "rb") as f:
-            self.nameToModuleHash.update(
-                SerializationContext().deserialize(f.read(), Dict(str, str))
-            )
+            func_name_to_module_hash = SerializationContext().deserialize(f.read(), Dict(str, str))
+
+            for func_name, module_hash in func_name_to_module_hash.items():
+                link_name = self._generate_link_name(func_name, module_hash)
+                self.func_name_to_link_names.setdefault(func_name, []).append(link_name)
+                self.link_name_to_module_hash[link_name] = module_hash
 
         self.moduleManifestsLoaded.add(moduleHash)
 
-    def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodules, dependencyEdgelist):
+    def writeModuleToDisk(self, binarySharedObject, hashToUse, nameToTypedCallTarget, submodules, dependencyEdgelist):
         """Write out a disk representation of this module.
 
         This includes writing both the shared object, a manifest of the function names
@@ -235,7 +268,6 @@ def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodule
         to interact with the compiler cache simultaneously without relying on
         individual file-level locking.
         """
-        hashToUse = SerializationContext().sha_hash(str(uuid.uuid4())).hexdigest
 
         targetDir = os.path.join(
             self.cacheDir,
@@ -264,23 +296,20 @@ def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodule
             for sourceName in manifest:
                 f.write(sourceName + "\n")
 
-        # write the type manifest
         with open(os.path.join(tempTargetDir, "type_manifest.dat"), "wb") as f:
             f.write(SerializationContext().serialize(nameToTypedCallTarget))
 
-        # write the nativetype manifest
         with open(os.path.join(tempTargetDir, "native_type_manifest.dat"), "wb") as f:
             f.write(SerializationContext().serialize(binarySharedObject.functionTypes))
 
-        # write the type manifest
         with open(os.path.join(tempTargetDir, "globals_manifest.dat"), "wb") as f:
             f.write(SerializationContext().serialize(binarySharedObject.serializedGlobalVariableDefinitions))
 
         with open(os.path.join(tempTargetDir, "submodules.dat"), "wb") as f:
             f.write(SerializationContext().serialize(ListOf(str)(submodules), ListOf(str)))
 
         with open(os.path.join(tempTargetDir, "function_dependencies.dat"), "wb") as f:
-            f.write(SerializationContext().serialize(dependencyEdgelist))  # might need a listof
+            f.write(SerializationContext().serialize(dependencyEdgelist))
 
         with open(os.path.join(tempTargetDir, "global_dependencies.dat"), "wb") as f:
             f.write(SerializationContext().serialize(binarySharedObject.globalDependencies))
@@ -293,14 +322,15 @@ def writeModuleToDisk(self, binarySharedObject, nameToTypedCallTarget, submodule
             else:
                 shutil.rmtree(tempTargetDir)
 
-        return targetDir, hashToUse
+        return targetDir
 
-    def function_pointer_by_name(self, linkName):
-        moduleHash = self.nameToModuleHash.get(linkName)
+    def function_pointer_by_name(self, func_name):
+        linkName = self._select_link_name(func_name)
+        moduleHash = self.link_name_to_module_hash.get(linkName)
         if moduleHash is None:
             raise Exception("Can't find a module for " + linkName)
 
         if moduleHash not in self.loadedBinarySharedObjects:
             self.loadForSymbol(linkName)
 
-        return self.loadedBinarySharedObjects[moduleHash].functionPointers[linkName]
+        return self.loadedBinarySharedObjects[moduleHash].functionPointers[func_name]