Account for memory allocated by trino s3 staging output stream

[homar]

Description

Account for memory allocated by TrinoS3FileSystems
Fixes: #14023

Non-technical explanation

Some queries that previously killed cluster should be gracefully terminated now

Release notes

( ) This is not user-visible or docs only and no release notes are required.
( ) Release notes are required, please propose a release note for me.
(x) Release notes are required, with the following suggested text:

# Core
* Account for memory allocated by TrinoS3FileSystems. ({issue}`14023`)

[sopel39]

getOutputStreamRetainedSizeInBytesSupplier -> retainedSizeInBytesSupplier

[sopel39]

you can just extend io.trino.hdfs.TrinoFileSystemCache.OutputStreamWrapper

[sopel39]

@findepi @electrum why do we need OutputStreamWrapper anyway?

Seems like io.trino.hdfs.TrinoFileSystemCache.OutputStreamWrapper#getWrappedStream

        @Override
        public OutputStream getWrappedStream()
        {
            return ((FSDataOutputStream) super.getWrappedStream()).getWrappedStream();
        }

is fragile assumption

[homar]

you can just extend io.trino.hdfs.TrinoFileSystemCache.OutputStreamWrapper

it is private i didn't want to change that, but sure I can do it

[sopel39]

it is private i didn't want to change that, but sure I can do it

Oh, but I mean you can make OutputStreamWrapper do what MemoryAwareOutputStreamWrapper does (provide MemoryAware interface for wrapper stream)

[homar]

oh, right

[sopel39]

getOutputStreamRetainedSizeInBytesSupplier -> outputStreamRetainedSizeInBytesSupplier

[sopel39]

If transferManager does not hold for memory outside between TrinoS3StagingOutputStream calls, then it should be 0

[homar]

I was not sure about this but you are right probably

[sopel39]

keep it private

[sopel39]

getOutputStreamRetainedSizeInBytesSupplier -> outputStreamRetainedSizeInBytesSupplier

[sopel39]

I think clearer approach is to use io.trino.filesystem.TrinoOutputFile in writers and modify io.trino.filesystem.TrinoOutputFile#create:

    OutputStream create(AggregatedMemoryContext context);

    OutputStream createOrOverwrite(AggregatedMemoryContext context)

then in io.trino.filesystem.hdfs.HdfsOutputFile#create(boolean) you can just check if FileSystem is TrinoS3FileSystem (see io.trino.plugin.hive.util.HiveWriteUtils#isS3FileSystem) and use dedicated create method (to be created) in TrinoS3FileSystem.

This way there is less intermediate layers and OutputStream unwrapping. It should be easy to test (just check if TrinoOutputFile accounts for memory in S3) and it won't be that fragile

[homar]

in io.trino.filesystem.hdfs.HdfsOutputFile#create(boolean) you can just check if FileSystem is TrinoS3FileSystem

I can't because TrinoS3FileSystem belongs to trino-hive module that depends on trino-filesystem to which HdfsOutputFile belongs. So HdfsOutputFile can't see TrinoS3FileSystem so there is no way to push AggregatedMemoryContext down

[sopel39]

I can't because TrinoS3FileSystem belongs to trino-hive module that depends on trino-filesystem to which HdfsOutputFile belongs. So HdfsOutputFile can't see TrinoS3FileSystem so there is no way to push AggregatedMemoryContext down

You can add MemoryAwareFileSystem interface and make TrinoS3FileSystem extend it.
The point is that such approach would be testable (e.g. you can add a test that io.trino.filesystem.hdfs.HdfsOutputFile with S3 path is actually memory aware)

[homar]

I can't because TrinoS3FileSystem belongs to trino-hive module that depends on trino-filesystem to which HdfsOutputFile belongs. So HdfsOutputFile can't see TrinoS3FileSystem so there is no way to push AggregatedMemoryContext down

You can add MemoryAwareFileSystem interface and make TrinoS3FileSystem extend it. The point is that such approach would be testable (e.g. you can add a test that io.trino.filesystem.hdfs.HdfsOutputFile with S3 path is actually memory aware)

That is not the only problem. Those create methods from TrinoOutputFile are used not only by writers but also in other places like io.trino.filesystem.fileio.ForwardingOutputFile. And in some of these places, like the one mentioned or S3TransactionLogSynchronizer it is not really feasible to obtain instance of AggregatedMemoryContext to pass it down. I know I can create them in place using newRootAggregatedMemoryContext but then nobody will read them in some execution paths so it will look badly imho. Or am I missing something like only adding these 2 methods mentioned by you so the TrinoOutputFile interface would like:

 OutputStream create(AggregatedMemoryContext context);

 OutputStream createOrOverwrite(AggregatedMemoryContext context)

 OutputStream create();

 OutputStream createOrOverwrite();

and simply creating instances of AggregatedMemoryContext for methods that don't accept them - though they theirs values won't be ever read

[homar]

@sopel39 please take another look.
There are some things I am not sure about - like what I did in

    public FSDataOutputStream create(Path path, FsPermission permission, boolean overwrite, int bufferSize, short replication, long blockSize, Progressable progress)

in TrinoS3FileSystem

[sopel39]

private

[sopel39]

AggregatedMemoryContext aggregatedMemoryContext -> AggregatedMemoryContext memoryContext

[sopel39]

AggregatedMemoryContext aggregatedMemoryContext -> AggregatedMemoryContext memoryContext

[sopel39]

AggregatedMemoryContext aggregatedMemoryContext -> AggregatedMemoryContext memoryContext

[sopel39]

aggregatedMemoryContext -> outputMemoryContext

[sopel39]

make this constructor private.

Add factory methods:

OutputStreamOrcDataSink OutputStreamOrcDataSink#create(TrinoOutputFile outputFile) {
  AggregatedMemoryContext context = newSimpleAggregatedMemoryContext();
  return new OutputStreamOrcDataSink(outputFile.create(context), context);
}

/**
 * Use only when output stream memory consumption tracking is not possible.
 */
OutputStreamOrcDataSink OutputStreamOrcDataSink#create(OutputStream stream) {
  return new OutputStreamOrcDataSink(strea, newSimpleAggregatedMemoryContext());
}

[sopel39]

not needed

[sopel39]

nit: buffer.length -> 0

[sopel39]

make it TrinoOutputFile outputFile = fileSystem.newOutputFile(outputPath)

[sopel39]

static constructor above OutputStreamOrcDataSink constructor

[sopel39]

remove? you don't account for mem in TrinoS3StagingOutputStream

[sopel39]

buffer is already created above

[sopel39]

you don't really have to mock these. Use real objects (e.g io.trino.memory.context.AggregatedMemoryContext#newRootAggregatedMemoryContext with reservation handler to track max allocation

[sopel39]

similar as for OutputStreamOrcDataSink. Create AggregatedMemoryContext locally in ParquetFileWriter.
You can then simplify calls to createParquetFileWriter

[sopel39]

create memory context locally and override getMemoryUsage method

[sopel39]

why reserved is not 0? Mem leak?

[homar]

Not a memory leak. in the test I simply don't close outputStream, memory context is not closed and thus not set to 0.
I did it on purpose to be able to catche the value. But I can do it better.

[sopel39]

close outputStream too

[homar]

heh I didn't notice this before answering the comment above ;) it is done now

[sopel39]

lgtm % comments

[sopel39]

remove this dependency

[sopel39]

just close is sufficient (no need for setBytes)

[sopel39]

I would just store max and make sure that assertThat(memoryReservationHandler.getMaxReserver).isGreatherThan(0);

Checking that 64 is part of reserved values is fragile (depends on implementation)