sjvanrossum commented on PR #39458:
URL: https://github.com/apache/beam/pull/39458#issuecomment-5118934942
Something like this perhaps?
I'll run TSAN checks in a bit.
```
/** The channel to write to. */
private volatile @Nullable WritableByteChannel channel;
/**
* Opens a uniquely named temporary file and initializes the writer
using {@link #prepareWrite}.
*
* <p>The unique id that is given to open should be used to ensure that
the writer's output does
* not interfere with the output of other Writers, as a bundle may be
executed many times for
* fault tolerance.
*/
public final void open(String uId) throws Exception {
// Shadow members as local variables to ensure local consistency and
force member access
// through this.
final @Nullable String id;
final @Nullable ResourceId outputFile;
final @Nullable WritableByteChannel channel;
try {
this.id = id = spreadUid(uId); // normal store
ResourceId tempDirectory = getWriteOperation().getTempDirectory();
this.outputFile =
outputFile =
tempDirectory.resolve(id,
StandardResolveOptions.RESOLVE_FILE); // normal store
verifyNotNull(
outputFile,
"FileSystems are not allowed to return null from resolve: %s",
tempDirectory);
final WritableByteChannelFactory factory =
getWriteOperation().getSink().writableByteChannelFactory;
// The factory may force a MIME type or it may return null,
indicating to use the sink's
// MIME.
String channelMimeType = firstNonNull(factory.getMimeType(),
mimeType);
CreateOptions createOptions =
StandardCreateOptions.builder()
.setMimeType(channelMimeType)
// The file is based upon a uuid and thus we expect it to be
unique and to not
// already
// exist. A new uuid is generated on each bundle processing
and thus this also holds
// across bundle retries. Collisions of filenames would
result in data loss as we
// would otherwise overwrite already finalized data.
.setExpectFileToNotExist(true)
.build();
WritableByteChannel tempChannel = FileSystems.create(outputFile,
createOptions);
try {
this.channel =
channel =
factory.create(tempChannel); // volatile store
(sequentially consistent release)
} catch (Exception e) {
// If we have opened the underlying channel but fail to open the
compression channel,
// we should still close the underlying channel.
try (tempChannel) {
throw e;
}
}
} catch (Throwable t) {
// Ensure that memory visibility matches non-exception case before
throwing.
this.channel = null; // volatile store (sequentially consistent
release)
throw t;
}
// The caller shouldn't have to close() this Writer if it fails to
open(), so close
// the channel if prepareWrite() or writeHeader() fails.
try {
LOG.debug("Preparing write to {}.", outputFile);
prepareWrite(channel);
LOG.debug("Writing header to {}.", outputFile);
writeHeader();
} catch (Exception e) {
LOG.error("Beginning write to {} failed, closing channel.",
outputFile, e);
try (channel) {
throw e;
}
}
LOG.debug("Starting write of bundle {} to {}.", id, outputFile);
}
public final void cleanup() throws Exception {
final @Nullable WritableByteChannel channel = this.channel; //
volatile load (acquire)
final @Nullable ResourceId outputFile = this.outputFile; // normal load
try (channel) {
if (channel != null && channel.isOpen()) {
LOG.warn("Channel to {} is still open after cleanup() was
called.", outputFile);
}
// channel and outputFile may be null if open() was not called or
failed.
if (outputFile != null) {
LOG.info("Deleting temporary file {}", outputFile);
FileSystems.delete(
Collections.singletonList(outputFile),
StandardMoveOptions.IGNORE_MISSING_FILES);
}
}
}
/** Closes the channel and returns the bundle result. */
public final void close() throws Exception {
final @Nullable WritableByteChannel channel = this.channel; //
volatile load (acquire)
final @Nullable ResourceId outputFile = this.outputFile; // normal load
try (channel) {
checkState(outputFile != null, "FileResult.close cannot be called
with a null outputFile");
LOG.debug("Closing {}", outputFile);
writeFooter();
finishWrite();
LOG.debug("Closing channel to {}.", outputFile);
}
LOG.info("Successfully wrote temporary file {}", outputFile);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
