Skip to content

Commit

Permalink
Remove DownloadManager instance from RegistryFactoryImpl.
Browse files Browse the repository at this point in the history 
This PR removes the DownloadManager from the registry factory implementation.
As the registry created by the factory is cached by the SkyFunction mechanism,
the DownloadManager instance was living too long - it is supposed to be
re-created for every command instantiation to respect changes in command line options,
but for the registry, it ignored those changes.

Instead, the DownloadManager is set directly into the affected SkyFunctions that
require access to it. This way, the per-command DownloadManager instance is
correctly used.
  • Loading branch information
@criemen
criemen committed Nov 5, 2024
1 parent 540958d commit 5636829
Show file tree
Hide file tree
Showing 8 changed files with 69 additions and 49 deletions.
23 changes: 14 additions & 9 deletions src/main/java/com/google/devtools/build/lib/bazel/BazelRepositoryModule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,9 @@ public class BazelRepositoryModule extends BlazeModule {
private List<String> allowedYankedVersions = ImmutableList.of();
private boolean disableNativeRepoRules;
private SingleExtensionEvalFunction singleExtensionEvalFunction;
private ModuleFileFunction moduleFileFunction;
private RepoSpecFunction repoSpecFunction;
private YankedVersionsFunction yankedVersionsFunction;

private final VendorCommand vendorCommand = new VendorCommand(clientEnvironmentSupplier);
private final RegistryFactoryImpl registryFactory =
Expand Down Expand Up @@ -248,14 +251,14 @@ public void workspaceInit(
builtinModules = ModuleFileFunction.getBuiltinModules(directories.getEmbeddedBinariesRoot());
}

moduleFileFunction =
new ModuleFileFunction(runtime.getRuleClassProvider().getBazelStarlarkEnvironment(), directories.getWorkspace(), builtinModules);
repoSpecFunction = new RepoSpecFunction();
yankedVersionsFunction = new YankedVersionsFunction();

builder
.addSkyFunction(SkyFunctions.REPOSITORY_DIRECTORY, repositoryDelegatorFunction)
.addSkyFunction(
SkyFunctions.MODULE_FILE,
new ModuleFileFunction(
runtime.getRuleClassProvider().getBazelStarlarkEnvironment(),
directories.getWorkspace(),
builtinModules))
.addSkyFunction(SkyFunctions.MODULE_FILE, moduleFileFunction)
.addSkyFunction(SkyFunctions.BAZEL_DEP_GRAPH, new BazelDepGraphFunction())
.addSkyFunction(
SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace()))
Expand All @@ -269,8 +272,8 @@ SkyFunctions.BAZEL_LOCK_FILE, new BazelLockFileFunction(directories.getWorkspace
.addSkyFunction(
SkyFunctions.REGISTRY,
new RegistryFunction(registryFactory, directories.getWorkspace()))
.addSkyFunction(SkyFunctions.REPO_SPEC, new RepoSpecFunction())
.addSkyFunction(SkyFunctions.YANKED_VERSIONS, new YankedVersionsFunction())
.addSkyFunction(SkyFunctions.REPO_SPEC, repoSpecFunction)
.addSkyFunction(SkyFunctions.YANKED_VERSIONS, yankedVersionsFunction)
.addSkyFunction(
SkyFunctions.VENDOR_FILE,
new VendorFileFunction(runtime.getRuleClassProvider().getBazelStarlarkEnvironment()))
Expand Down Expand Up @@ -305,8 +308,10 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
DownloadManager downloadManager =
new DownloadManager(repositoryCache, env.getDownloaderDelegate(), env.getHttpDownloader());
this.starlarkRepositoryFunction.setDownloadManager(downloadManager);
this.moduleFileFunction.setDownloadManager(downloadManager);
this.repoSpecFunction.setDownloadManager(downloadManager);
this.yankedVersionsFunction.setDownloadManager(downloadManager);
this.vendorCommand.setDownloadManager(downloadManager);
this.registryFactory.setDownloadManager(downloadManager);

clientEnvironmentSupplier.set(env.getRepoEnv());
PackageOptions pkgOptions = env.getOptions().getOptions(PackageOptions.class);
Expand Down
35 changes: 17 additions & 18 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/IndexRegistry.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,6 @@ public enum KnownFileHashesMode {
}

private final URI uri;
private final DownloadManager downloadManager;
private final Map<String, String> clientEnv;
private final Gson gson;
private final ImmutableMap<String, Optional<Checksum>> knownFileHashes;
Expand All @@ -99,14 +98,12 @@ public enum KnownFileHashesMode {

public IndexRegistry(
URI uri,
DownloadManager downloadManager,
Map<String, String> clientEnv,
ImmutableMap<String, Optional<Checksum>> knownFileHashes,
KnownFileHashesMode knownFileHashesMode,
ImmutableMap<ModuleKey, String> previouslySelectedYankedVersions,
Optional<Path> vendorDir) {
this.uri = uri;
this.downloadManager = downloadManager;
this.clientEnv = clientEnv;
this.gson =
new GsonBuilder()
Expand Down Expand Up @@ -136,9 +133,9 @@ private String constructUrl(String base, String... segments) {

/** Grabs a file from the given URL. Returns {@link Optional#empty} if the file doesn't exist. */
private Optional<byte[]> grabFile(
String url, ExtendedEventHandler eventHandler, boolean useChecksum)
String url, ExtendedEventHandler eventHandler, DownloadManager downloadManager, boolean useChecksum)
throws IOException, InterruptedException {
var maybeContent = doGrabFile(url, eventHandler, useChecksum);
var maybeContent = doGrabFile(downloadManager, url, eventHandler, useChecksum);
if ((knownFileHashesMode == KnownFileHashesMode.USE_AND_UPDATE
|| knownFileHashesMode == KnownFileHashesMode.USE_IMMUTABLE_AND_UPDATE)
&& useChecksum) {
Expand All @@ -148,7 +145,7 @@ private Optional<byte[]> grabFile(
}

private Optional<byte[]> doGrabFile(
String rawUrl, ExtendedEventHandler eventHandler, boolean useChecksum)
DownloadManager downloadManager, String rawUrl, ExtendedEventHandler eventHandler, boolean useChecksum)
throws IOException, InterruptedException {
Optional<Checksum> checksum;
if (knownFileHashesMode != KnownFileHashesMode.IGNORE && useChecksum) {
Expand Down Expand Up @@ -225,11 +222,11 @@ private Optional<byte[]> doGrabFile(
}

@Override
public Optional<ModuleFile> getModuleFile(ModuleKey key, ExtendedEventHandler eventHandler)
public Optional<ModuleFile> getModuleFile(ModuleKey key, ExtendedEventHandler eventHandler, DownloadManager downloadManager)
throws IOException, InterruptedException {
String url =
constructUrl(getUrl(), "modules", key.name(), key.version().toString(), "MODULE.bazel");
Optional<byte[]> maybeContent = grabFile(url, eventHandler, /* useChecksum= */ true);
Optional<byte[]> maybeContent = grabFile(url, eventHandler, downloadManager, /* useChecksum= */ true);
return maybeContent.map(content -> ModuleFile.create(content, url));
}

Expand Down Expand Up @@ -276,19 +273,19 @@ private static class GitRepoSourceJson {
* if the file doesn't exist.
*/
private Optional<String> grabJsonFile(
String url, ExtendedEventHandler eventHandler, boolean useChecksum)
String url, ExtendedEventHandler eventHandler, DownloadManager downloadManager, boolean useChecksum)
throws IOException, InterruptedException {
return grabFile(url, eventHandler, useChecksum).map(value -> new String(value, UTF_8));
return grabFile(url, eventHandler, downloadManager, useChecksum).map(value -> new String(value, UTF_8));
}

/**
* Grabs a JSON file from the given URL, and returns it as a parsed object with fields in {@code
* T}. Returns {@link Optional#empty} if the file doesn't exist.
*/
private <T> Optional<T> grabJson(
String url, Class<T> klass, ExtendedEventHandler eventHandler, boolean useChecksum)
String url, Class<T> klass, ExtendedEventHandler eventHandler, DownloadManager downloadManager, boolean useChecksum)
throws IOException, InterruptedException {
Optional<String> jsonString = grabJsonFile(url, eventHandler, useChecksum);
Optional<String> jsonString = grabJsonFile(url, eventHandler, downloadManager, useChecksum);
if (jsonString.isEmpty() || jsonString.get().isBlank()) {
return Optional.empty();
}
Expand All @@ -306,10 +303,10 @@ private <T> T parseJson(String jsonString, String url, Class<T> klass) throws IO
}

@Override
public RepoSpec getRepoSpec(ModuleKey key, ExtendedEventHandler eventHandler)
public RepoSpec getRepoSpec(ModuleKey key, ExtendedEventHandler eventHandler, DownloadManager downloadManager)
throws IOException, InterruptedException {
String jsonUrl = getSourceJsonUrl(key);
Optional<String> jsonString = grabJsonFile(jsonUrl, eventHandler, /* useChecksum= */ true);
Optional<String> jsonString = grabJsonFile(jsonUrl, eventHandler, downloadManager, /* useChecksum= */ true);
if (jsonString.isEmpty()) {
throw new FileNotFoundException(
String.format(
Expand All @@ -320,12 +317,12 @@ public RepoSpec getRepoSpec(ModuleKey key, ExtendedEventHandler eventHandler)
case "archive" -> {
ArchiveSourceJson typedSourceJson =
parseJson(jsonString.get(), jsonUrl, ArchiveSourceJson.class);
return createArchiveRepoSpec(typedSourceJson, getBazelRegistryJson(eventHandler), key);
return createArchiveRepoSpec(typedSourceJson, getBazelRegistryJson(eventHandler, downloadManager), key);
}
case "local_path" -> {
LocalPathSourceJson typedSourceJson =
parseJson(jsonString.get(), jsonUrl, LocalPathSourceJson.class);
return createLocalPathRepoSpec(typedSourceJson, getBazelRegistryJson(eventHandler), key);
return createLocalPathRepoSpec(typedSourceJson, getBazelRegistryJson(eventHandler, downloadManager), key);
}
case "git_repository" -> {
GitRepoSourceJson typedSourceJson =
Expand All @@ -343,7 +340,7 @@ private String getSourceJsonUrl(ModuleKey key) {
getUrl(), "modules", key.name(), key.version().toString(), SOURCE_JSON_FILENAME);
}

private Optional<BazelRegistryJson> getBazelRegistryJson(ExtendedEventHandler eventHandler)
private Optional<BazelRegistryJson> getBazelRegistryJson(ExtendedEventHandler eventHandler, DownloadManager downloadManager)
throws IOException, InterruptedException {
if (bazelRegistryJson == null || bazelRegistryJsonEvents == null) {
synchronized (this) {
Expand All @@ -355,6 +352,7 @@ private Optional<BazelRegistryJson> getBazelRegistryJson(ExtendedEventHandler ev
constructUrl(getUrl(), "bazel_registry.json"),
BazelRegistryJson.class,
storedEventHandler,
downloadManager,
/* useChecksum= */ true);
bazelRegistryJsonEvents = storedEventHandler;
}
Expand Down Expand Up @@ -484,13 +482,14 @@ private RepoSpec createGitRepoSpec(GitRepoSourceJson sourceJson) {

@Override
public Optional<ImmutableMap<Version, String>> getYankedVersions(
String moduleName, ExtendedEventHandler eventHandler)
String moduleName, ExtendedEventHandler eventHandler, DownloadManager downloadManager)
throws IOException, InterruptedException {
Optional<MetadataJson> metadataJson =
grabJson(
constructUrl(getUrl(), "modules", moduleName, "metadata.json"),
MetadataJson.class,
eventHandler,
downloadManager,
// metadata.json is not immutable
/* useChecksum= */ false);
if (metadataJson.isEmpty()) {
Expand Down
8 changes: 7 additions & 1 deletion src/main/java/com/google/devtools/build/lib/bazel/bzlmod/ModuleFileFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import com.google.devtools.build.lib.bazel.bzlmod.ModuleFileValue.RootModuleFileValue;
import com.google.devtools.build.lib.bazel.repository.PatchUtil;
import com.google.devtools.build.lib.bazel.repository.downloader.Checksum.MissingChecksumException;
import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
import com.google.devtools.build.lib.cmdline.Label;
import com.google.devtools.build.lib.cmdline.LabelConstants;
import com.google.devtools.build.lib.cmdline.LabelSyntaxException;
Expand Down Expand Up @@ -99,6 +100,7 @@ public class ModuleFileFunction implements SkyFunction {
private final BazelStarlarkEnvironment starlarkEnv;
private final Path workspaceRoot;
private final ImmutableMap<String, NonRegistryOverride> builtinModules;
@Nullable private DownloadManager downloadManager;

private static final String BZLMOD_REMINDER =
"""
Expand Down Expand Up @@ -231,6 +233,10 @@ public SkyValue compute(SkyKey skyKey, Environment env)
getModuleFileResult.downloadEventHandler.getPosts()));
}

public void setDownloadManager(DownloadManager downloadManager) {
this.downloadManager = downloadManager;
}

@Nullable
private SkyValue computeForRootModule(
StarlarkSemantics starlarkSemantics, Environment env, SymbolGenerator<?> symbolGenerator)
Expand Down Expand Up @@ -620,7 +626,7 @@ private GetModuleFileResult getModuleFile(
StoredEventHandler downloadEventHandler = new StoredEventHandler();
for (Registry registry : registryObjects) {
try {
Optional<ModuleFile> maybeModuleFile = registry.getModuleFile(key, downloadEventHandler);
Optional<ModuleFile> maybeModuleFile = registry.getModuleFile(key, downloadEventHandler, this.downloadManager);
if (maybeModuleFile.isEmpty()) {
continue;
}
Expand Down
7 changes: 4 additions & 3 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/Registry.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package com.google.devtools.build.lib.bazel.bzlmod;

import com.google.common.collect.ImmutableMap;
import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
import com.google.devtools.build.lib.events.ExtendedEventHandler;
import com.google.devtools.build.skyframe.NotComparableSkyValue;
import java.io.IOException;
Expand All @@ -31,22 +32,22 @@ public interface Registry extends NotComparableSkyValue {
* Retrieves the contents of the module file of the module identified by {@code key} from the
* registry. Returns {@code Optional.empty()} when the module is not found in this registry.
*/
Optional<ModuleFile> getModuleFile(ModuleKey key, ExtendedEventHandler eventHandler)
Optional<ModuleFile> getModuleFile(ModuleKey key, ExtendedEventHandler eventHandler, DownloadManager downloadManager)
throws IOException, InterruptedException;

/**
* Retrieves the {@link RepoSpec} object that indicates how the contents of the module identified
* by {@code key} should be materialized as a repo.
*/
RepoSpec getRepoSpec(ModuleKey key, ExtendedEventHandler eventHandler)
RepoSpec getRepoSpec(ModuleKey key, ExtendedEventHandler eventHandler, DownloadManager downloadManager)
throws IOException, InterruptedException;

/**
* Retrieves yanked versions of the module identified by {@code key.getName()} from the registry.
* Returns {@code Optional.empty()} when the information is not found in the registry.
*/
Optional<ImmutableMap<Version, String>> getYankedVersions(
String moduleName, ExtendedEventHandler eventHandler)
String moduleName, ExtendedEventHandler eventHandler, DownloadManager downloadManager)
throws IOException, InterruptedException;

/**
Expand Down
7 changes: 0 additions & 7 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RegistryFactoryImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import com.google.devtools.build.lib.bazel.bzlmod.IndexRegistry.KnownFileHashesMode;
import com.google.devtools.build.lib.bazel.repository.RepositoryOptions.LockfileMode;
import com.google.devtools.build.lib.bazel.repository.downloader.Checksum;
import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
import com.google.devtools.build.lib.vfs.Path;
import java.net.URI;
import java.net.URISyntaxException;
Expand All @@ -30,17 +29,12 @@

/** Prod implementation of {@link RegistryFactory}. */
public class RegistryFactoryImpl implements RegistryFactory {
@Nullable private DownloadManager downloadManager;
private final Supplier<Map<String, String>> clientEnvironmentSupplier;

public RegistryFactoryImpl(Supplier<Map<String, String>> clientEnvironmentSupplier) {
this.clientEnvironmentSupplier = clientEnvironmentSupplier;
}

public void setDownloadManager(DownloadManager downloadManager) {
this.downloadManager = downloadManager;
}

@Override
public Registry createRegistry(
String url,
Expand Down Expand Up @@ -76,7 +70,6 @@ public Registry createRegistry(
};
return new IndexRegistry(
uri,
downloadManager,
clientEnvironmentSupplier.get(),
knownFileHashes,
knownFileHashesMode,
Expand Down
9 changes: 7 additions & 2 deletions src/main/java/com/google/devtools/build/lib/bazel/bzlmod/RepoSpecFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
//

package com.google.devtools.build.lib.bazel.bzlmod;

import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
import com.google.devtools.build.lib.events.StoredEventHandler;
import com.google.devtools.build.lib.profiler.Profiler;
import com.google.devtools.build.lib.profiler.ProfilerTask;
Expand All @@ -32,6 +32,7 @@
* fetching required information from its {@link Registry}.
*/
public class RepoSpecFunction implements SkyFunction {
@Nullable private DownloadManager downloadManager;

@Override
@Nullable
Expand All @@ -49,7 +50,7 @@ public SkyValue compute(SkyKey skyKey, Environment env)
try (SilentCloseable c =
Profiler.instance()
.profile(ProfilerTask.BZLMOD, () -> "compute repo spec: " + key.getModuleKey())) {
repoSpec = registry.getRepoSpec(key.getModuleKey(), downloadEvents);
repoSpec = registry.getRepoSpec(key.getModuleKey(), downloadEvents, this.downloadManager);
} catch (IOException e) {
throw new RepoSpecException(
ExternalDepsException.withCauseAndMessage(
Expand All @@ -63,6 +64,10 @@ public SkyValue compute(SkyKey skyKey, Environment env)
repoSpec, RegistryFileDownloadEvent.collectToMap(downloadEvents.getPosts()));
}

public void setDownloadManager(DownloadManager downloadManager) {
this.downloadManager = downloadManager;
}

static final class RepoSpecException extends SkyFunctionException {

RepoSpecException(ExternalDepsException cause) {
Expand Down
8 changes: 7 additions & 1 deletion src/main/java/com/google/devtools/build/lib/bazel/bzlmod/YankedVersionsFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

package com.google.devtools.build.lib.bazel.bzlmod;

import com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager;
import com.google.devtools.build.lib.events.Event;
import com.google.devtools.build.lib.profiler.Profiler;
import com.google.devtools.build.lib.profiler.ProfilerTask;
Expand All @@ -31,6 +32,7 @@
* Registry}.
*/
public class YankedVersionsFunction implements SkyFunction {
@Nullable private DownloadManager downloadManager;

@Override
@Nullable
Expand All @@ -47,7 +49,7 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
.profile(
ProfilerTask.BZLMOD, () -> "getting yanked versions: " + key.getModuleName())) {
return YankedVersionsValue.create(
registry.getYankedVersions(key.getModuleName(), env.getListener()));
registry.getYankedVersions(key.getModuleName(), env.getListener(), downloadManager));
} catch (IOException e) {
env.getListener()
.handle(
Expand All @@ -60,4 +62,8 @@ public SkyValue compute(SkyKey skyKey, Environment env) throws InterruptedExcept
return YankedVersionsValue.create(Optional.empty());
}
}

public void setDownloadManager(DownloadManager downloadManager) {
this.downloadManager = downloadManager;
}
}
Loading

0 comments on commit 5636829

Please sign in to comment.