6060import java .util .ArrayList ;
6161import java .util .Arrays ;
6262import java .util .Deque ;
63- import java .util .HashSet ;
6463import java .util .List ;
6564import java .util .Set ;
6665import java .util .concurrent .ConcurrentHashMap ;
@@ -87,53 +86,124 @@ public abstract class AbstractActionInputPrefetcher implements ActionInputPrefet
8786
8887 private final Set <ActionInput > missingActionInputs = Sets .newConcurrentHashSet ();
8988
90- // Tracks the number of ongoing prefetcher calls temporarily making an output directory writable.
91- // Since concurrent calls may write to the same directory, it's not safe to make it non-writable
92- // until no other ongoing calls are writing to it.
93- private final ConcurrentHashMap <Path , Integer > temporarilyWritableDirectories =
89+ private static final Object dummyValue = new Object ();
90+
91+ /**
92+ * Tracks output directories temporarily made writable for prefetching. Since concurrent calls may
93+ * write to the same directory, it's not safe to make it non-writable until no other ongoing
94+ * prefetcher calls are writing to it.
95+ */
96+ private final ConcurrentHashMap <Path , DirectoryState > temporarilyWritableDirectories =
9497 new ConcurrentHashMap <>();
9598
96- /** Keeps track of output directories written to by a single prefetcher call. */
99+ /** The state of a single temporarily writable directory. */
100+ private static final class DirectoryState {
101+ /** The number of ongoing prefetcher calls touching this directory. */
102+ int numCalls ;
103+ /** Whether the output permissions must be set on the directory when prefetching completes. */
104+ boolean mustSetOutputPermissions ;
105+ }
106+
107+ /**
108+ * Tracks output directories written to by a single prefetcher call.
109+ *
110+ * <p>This makes it possible to set the output permissions on directories touched by the
111+ * prefetcher call all at once, so that files prefetched within the same call don't repeatedly set
112+ * output permissions on the same directory.
113+ */
97114 private final class DirectoryContext {
98- private final HashSet <Path > dirs = new HashSet <>();
115+ private final ConcurrentHashMap <Path , Object > dirs = new ConcurrentHashMap <>();
99116
100117 /**
101- * Adds to the set of directories written to by the prefetcher call associated with this
102- * context.
118+ * Makes a directory temporarily writable for the remainder of the prefetcher call associated
119+ * with this context.
120+ *
121+ * @param isDefinitelyTreeDir Whether this directory definitely belongs to a tree artifact.
122+ * Otherwise, whether it belongs to a tree artifact is inferred from its permissions.
103123 */
104- void add (Path dir ) {
105- if (dirs .add (dir )) {
106- temporarilyWritableDirectories .compute (dir , (unused , count ) -> count != null ? ++count : 1 );
124+ void createOrSetWritable (Path dir , boolean isDefinitelyTreeDir ) throws IOException {
125+ AtomicReference <IOException > caughtException = new AtomicReference <>();
126+
127+ dirs .compute (
128+ dir ,
129+ (outerUnused , previousValue ) -> {
130+ if (previousValue != null ) {
131+ return previousValue ;
132+ }
133+
134+ temporarilyWritableDirectories .compute (
135+ dir ,
136+ (innerUnused , state ) -> {
137+ if (state == null ) {
138+ state = new DirectoryState ();
139+ state .numCalls = 0 ;
140+
141+ try {
142+ if (isDefinitelyTreeDir ) {
143+ state .mustSetOutputPermissions = true ;
144+ var ignored = dir .createWritableDirectory ();
145+ } else {
146+ // If the directory is writable, it's a package and should be kept writable.
147+ // Otherwise, it must belong to a tree artifact, since the directory for a
148+ // tree is created in a non-writable state before prefetching begins, and
149+ // this is the first time the prefetcher is seeing it.
150+ state .mustSetOutputPermissions = !dir .isWritable ();
151+ if (state .mustSetOutputPermissions ) {
152+ dir .setWritable (true );
153+ }
154+ }
155+ } catch (IOException e ) {
156+ caughtException .set (e );
157+ return null ;
158+ }
159+ }
160+
161+ ++state .numCalls ;
162+
163+ return state ;
164+ });
165+
166+ if (caughtException .get () != null ) {
167+ return null ;
168+ }
169+
170+ return dummyValue ;
171+ });
172+
173+ if (caughtException .get () != null ) {
174+ throw caughtException .get ();
107175 }
108176 }
109177
110178 /**
111179 * Signals that the prefetcher call associated with this context has finished.
112180 *
113- * <p>The output permissions will be set on any directories written to by this call that are not
114- * being written to by other concurrent calls .
181+ * <p>The output permissions will be set on any directories temporarily made writable by this
182+ * call, if this is the last remaining call temporarily making them writable .
115183 */
116184 void close () throws IOException {
117185 AtomicReference <IOException > caughtException = new AtomicReference <>();
118186
119- for (Path dir : dirs ) {
187+ for (Path dir : dirs . keySet () ) {
120188 temporarilyWritableDirectories .compute (
121189 dir ,
122- (unused , count ) -> {
123- checkState (count != null );
124- if (--count == 0 ) {
125- try {
126- dir .chmod (outputPermissions .getPermissionsMode ());
127- } catch (IOException e ) {
128- // Store caught exceptions, but keep cleaning up the map.
129- if (caughtException .get () == null ) {
130- caughtException .set (e );
131- } else {
132- caughtException .get ().addSuppressed (e );
190+ (unused , state ) -> {
191+ checkState (state != null );
192+ if (--state .numCalls == 0 ) {
193+ if (state .mustSetOutputPermissions ) {
194+ try {
195+ dir .chmod (outputPermissions .getPermissionsMode ());
196+ } catch (IOException e ) {
197+ // Store caught exceptions, but keep cleaning up the map.
198+ if (caughtException .get () == null ) {
199+ caughtException .set (e );
200+ } else {
201+ caughtException .get ().addSuppressed (e );
202+ }
133203 }
134204 }
135205 }
136- return count > 0 ? count : null ;
206+ return state . numCalls > 0 ? state : null ;
137207 });
138208 }
139209 dirs .clear ();
@@ -518,24 +588,19 @@ private void finalizeDownload(
518588 }
519589 while (!dirs .isEmpty ()) {
520590 Path dir = dirs .pop ();
521- dirCtx .add (dir );
522591 // Create directory or make existing directory writable.
523- var unused = dir .createWritableDirectory ();
592+ // We know with certainty that the directory belongs to a tree artifact.
593+ dirCtx .createOrSetWritable (dir , /* isDefinitelyTreeDir= */ true );
524594 }
525595 } else {
526- // If the parent directory is not writable, temporarily make it so.
527- // This is needed when fetching a non-tree artifact nested inside a tree artifact, or a tree
528- // artifact inside a fileset (see b/254844173 for the latter).
529- // TODO(tjgq): Fix the TOCTTOU race between isWritable and setWritable. This requires keeping
530- // track of the original directory permissions. Note that nested artifacts are relatively rare
531- // and will eventually be disallowed (see issue #16729).
532- if (!parentDir .isWritable ()) {
533- dirCtx .add (parentDir );
534- parentDir .setWritable (true );
535- }
596+ // Temporarily make the parent directory writable if needed.
597+ // We don't know with certainty that the directory does not belong to a tree artifact; it
598+ // could if the fetched file is a non-tree artifact nested inside a tree artifact, or a
599+ // tree artifact inside a fileset (see b/254844173 for the latter).
600+ dirCtx .createOrSetWritable (parentDir , /* isDefinitelyTreeDir= */ false );
536601 }
537602
538- // Set output permissions on files (tree subdirectories are handled in stopPrefetching ),
603+ // Set output permissions on files (tree subdirectories are handled in DirectoryContext#close ),
539604 // matching the behavior of SkyframeActionExecutor#checkOutputs for artifacts produced by local
540605 // actions.
541606 tmpPath .chmod (outputPermissions .getPermissionsMode ());
0 commit comments