pipe-cli/mount/pipefuse/mpu.py

# Copyright 2017-2020 EPAM Systems, Inc. (https://www.epam.com/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import sys
from abc import abstractmethod, ABCMeta
from collections import namedtuple

import intervals

from pipefuse.fuseutils import MB, GB

_CopyPart = namedtuple('CopyPart', ['start', 'end', 'offset', 'part_number', 'part_path', 'keep'])


class UnmanageableMultipartUploadException(RuntimeError):

    def __init__(self, *args):
        super(UnmanageableMultipartUploadException, self).__init__(*args)


class MultipartUpload:
    __metaclass__ = ABCMeta

    @property
    @abstractmethod
    def path(self):
        pass

    @abstractmethod
    def initiate(self):
        pass

    @abstractmethod
    def upload_part(self, buf, offset=None, part_number=None, part_path=None, keep=False):
        pass

    @abstractmethod
    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        pass

    @abstractmethod
    def complete(self):
        pass

    @abstractmethod
    def abort(self):
        pass


class MultipartUploadDecorator(MultipartUpload):

    def __init__(self, mpu):
        self._mpu = mpu

    @property
    def path(self):
        return self._mpu.path

    def initiate(self):
        self._mpu.initiate()

    def upload_part(self, buf, offset=None, part_number=None, part_path=None, keep=False):
        self._mpu.upload_part(buf, offset, part_number)

    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        self._mpu.upload_copy_part(start, end, offset, part_number)

    def complete(self):
        self._mpu.complete()

    def abort(self):
        self._mpu.abort()


class _PartialChunk:
    __metaclass__ = ABCMeta

    @property
    @abstractmethod
    def offset(self):
        pass

    @abstractmethod
    def append(self, offset, buf):
        pass

    @abstractmethod
    def is_full(self):
        pass

    @abstractmethod
    def missing_intervals(self):
        pass

    @abstractmethod
    def collect(self):
        pass


class _IncompletePartialChunk(_PartialChunk):

    def __init__(self, offset, size):
        self._offset = offset
        self._size = size
        self._buf = bytearray(size)
        self._bounds_interval = intervals.closed(0, self._size)
        self._interval = intervals.empty()

    @property
    def offset(self):
        return self._offset

    def append(self, offset, buf):
        end = offset + len(buf)
        self._buf[offset:end] = buf[:]
        self._interval |= intervals.closed(offset, end)

    def is_full(self):
        return self._missing_interval().is_empty()

    def missing_intervals(self):
        for interval in self._missing_interval():
            yield interval.lower, interval.upper

    def _missing_interval(self):
        return self._interval.complement().intersection(self._bounds_interval)

    def collect(self):
        return self._buf[:self._interval.upper]


class _CompletePartialChunk(_PartialChunk):

    def __init__(self, offset):
        self._offset = offset

    @property
    def offset(self):
        return self._offset

    def append(self, offset, buf):
        pass

    def is_full(self):
        return True

    def missing_intervals(self):
        return []

    def collect(self):
        return bytearray(0)


class ChunkedMultipartUpload(MultipartUploadDecorator):

    def __init__(self, mpu, original_size, download, chunk_size, min_chunk, max_chunk):
        """
        Chunked multipart upload.

        Cuts all the incoming uploads into chunks of the given size.

        Has a limit on the maximum file size that can be written using chunked multipart upload.
        It can be calculated multiplying chunk size by max chunk number. F.e. for chunk size of 10MB and 10000 chunks
        it will be 100GB, for chunk size of 100MB and 10000 chunks it will be 1TB.

        Fills gaps between uploaded parts with the copy parts of unlimited size.

        :param mpu: Wrapping multipart upload.
        :param original_size: Destination file original size.
        :param download: Function that retrieves content from an object by its path, offset and length.
        :param chunk_size: Size of a single upload part.
        :param min_chunk: Minimum allowed chunk number.
        :param max_chunk: Maximum allowed chunk number.
        """
        super(ChunkedMultipartUpload, self).__init__(mpu)
        self._mpu = mpu
        self._original_size = original_size
        self._download = download
        self._chunk_size = chunk_size
        self._partial_chunks = {}
        self._min_chunk = min_chunk
        self._max_chunk = max_chunk

    def upload_part(self, buf, offset=None, part_number=None, part_path=None, keep=False):
        chunk = self._resolve_chunk(offset)
        chunk_offset = self._chunk_offset(chunk)
        chunk_shift = offset - chunk_offset
        buf_shift = 0
        while buf_shift < len(buf):
            if chunk_shift or buf_shift + self._chunk_size - chunk_shift > len(buf):
                partial_chunk = self._partial_chunks.get(chunk, None)
                if not partial_chunk:
                    partial_chunk = _IncompletePartialChunk(chunk_offset, self._chunk_size)
                    self._partial_chunks[chunk] = partial_chunk
                if partial_chunk.is_full():
                    raise UnmanageableMultipartUploadException(
                        'Multipart upload chunk %s cannot be reuploaded for %s.' % (chunk, self.path))
                partial_chunk.append(chunk_shift, buf[buf_shift:buf_shift + self._chunk_size - chunk_shift])
                if partial_chunk.is_full():
                    self._mpu.upload_part(partial_chunk.collect(), partial_chunk.offset, chunk)
                    self._partial_chunks[chunk] = _CompletePartialChunk(partial_chunk.offset)
            else:
                chunk_buf = bytearray(self._chunk_size)
                chunk_buf[chunk_shift:self._chunk_size] = buf[buf_shift:buf_shift + self._chunk_size]
                self._mpu.upload_part(chunk_buf, chunk_offset, chunk)
                self._partial_chunks[chunk] = _CompletePartialChunk(chunk_offset)
            buf_shift += self._chunk_size - chunk_shift
            chunk += 1
            chunk_offset += self._chunk_size
            chunk_shift = 0

    def _resolve_chunk(self, offset):
        first_chunk = self._min_chunk
        last_chunk = self._max_chunk
        while last_chunk - first_chunk > 1:
            mid_chunk = first_chunk + (last_chunk - first_chunk) / 2
            mid_chunk_offset = self._chunk_offset(mid_chunk)
            if offset > mid_chunk_offset:
                first_chunk = mid_chunk
            elif offset < mid_chunk_offset:
                last_chunk = mid_chunk
            else:
                return mid_chunk
        return last_chunk if offset >= self._chunk_offset(last_chunk) else first_chunk

    def _chunk_offset(self, chunk):
        return (chunk - 1) * self._chunk_size

    def complete(self):
        current_missing_chunk = 1
        for chunk in sorted(self._partial_chunks.keys()):
            if chunk > current_missing_chunk:
                missing_start = self._chunk_offset(current_missing_chunk)
                missing_end = self._chunk_offset(chunk)
                self._mpu.upload_copy_part(missing_start, missing_end, missing_start, current_missing_chunk)
            current_missing_chunk = chunk + 1
        last_written_chunk = max(self._partial_chunks.keys())
        last_missing_chunk = last_written_chunk + 1
        last_chunk_end = self._chunk_offset(last_missing_chunk)
        if last_chunk_end < self._original_size:
            self._mpu.upload_copy_part(last_chunk_end, self._original_size, last_chunk_end, last_missing_chunk)
        for chunk_number, partial_chunk in self._partial_chunks.items():
            if partial_chunk.is_full():
                continue
            for missing_start, missing_end in partial_chunk.missing_intervals():
                actual_start = missing_start + partial_chunk.offset
                actual_end = missing_end + partial_chunk.offset
                if self._original_size <= actual_start:
                    if chunk_number < last_written_chunk:
                        partial_chunk.append(missing_start,
                                             bytearray(actual_end - actual_start))
                elif actual_start < self._original_size < actual_end:
                    partial_chunk.append(missing_start,
                                         self._download(self.path, actual_start, self._original_size - actual_start))
                    if chunk_number < last_written_chunk:
                        partial_chunk.append(missing_start + self._original_size - actual_start,
                                             bytearray(actual_end - self._original_size))
                else:
                    partial_chunk.append(missing_start, self._download(self.path, actual_start, actual_end - actual_start))
            chunk = partial_chunk.collect()
            self._mpu.upload_part(chunk, partial_chunk.offset, chunk_number)
        self._mpu.complete()


class SplittingMultipartCopyUpload(MultipartUploadDecorator):

    def __init__(self, mpu, min_part_size=5 * MB, max_part_size=5 * GB):
        """
        Splitting multipart copy upload.

        Splits copy upload parts into several ones to fit maximum upload part size limit.
        Also takes into the account minimum upload part size.

        :param mpu: Wrapping multipart upload.
        :param min_part_size: Minimum upload part size.
        :param max_part_size: Maximum upload part size.
        """
        super(SplittingMultipartCopyUpload, self).__init__(mpu)
        self._mpu = mpu
        self._min_part_size = min_part_size
        self._max_part_size = max_part_size

    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        copy_part_length = end - start
        if copy_part_length > self._max_part_size:
            logging.debug('Splitting copy upload part %d into pieces for %s' % (part_number, self.path))
            remaining_length = copy_part_length
            current_offset = offset
            actual_part_number = part_number
            while remaining_length > 0:
                part_size = self._resolve_part_size(remaining_length)
                self._mpu.upload_copy_part(current_offset, current_offset + part_size, current_offset,
                                           actual_part_number)
                remaining_length -= part_size
                current_offset += part_size
                actual_part_number += 1
        else:
            self._mpu.upload_copy_part(start, end, offset, part_number)

    def _resolve_part_size(self, remaining_length):
        if self._min_part_size <= remaining_length <= self._max_part_size:
            return remaining_length
        else:
            return min(self._max_part_size, remaining_length - self._min_part_size)


class OutOfBoundsSplittingMultipartCopyUpload(SplittingMultipartCopyUpload):

    def __init__(self, mpu, original_size, min_part_size, max_part_size):
        """
        Out of bounds splitting multipart copy upload.

        Splits out of bounds copy upload parts into several ones to fit memory-safe part size limit.
        Also takes into the account minimum upload part size.

        :param mpu: Wrapping multipart upload.
        :param original_size: Destination file original size.
        :param min_part_size: Minimum upload part size.
        :param max_part_size: Maximum upload part size.
        """
        super(OutOfBoundsSplittingMultipartCopyUpload, self).__init__(mpu, min_part_size, max_part_size)
        self._original_size = original_size

    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        if self._original_size < end:
            super(OutOfBoundsSplittingMultipartCopyUpload, self).upload_copy_part(start, end, offset, part_number)
        else:
            self._mpu.upload_copy_part(start, end, offset, part_number)


class OutOfBoundsFillingMultipartCopyUpload(MultipartUploadDecorator):

    def __init__(self, mpu, original_size, download):
        """
        Out of bounds filling multipart copy upload.

        Fills copy upload part regions which are located beyond the original file size with null bytes.

        :param mpu: Wrapping multipart upload.
        :param original_size: Destination file original size.
        :param download: Function that retrieves content from the original file by its offset and length.
        """
        super(OutOfBoundsFillingMultipartCopyUpload, self).__init__(mpu)
        self._mpu = mpu
        self._original_size = original_size
        self._download = download

    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        if self._original_size <= start:
            logging.debug('Filling out of bounds copy upload part %s whole %d-%d for %s'
                          % (part_number, start, end, self.path))
            self._mpu.upload_part(bytearray(end - start), offset, part_number)
        elif start < self._original_size < end:
            logging.debug('Filling out of bounds copy upload part %s region %d-%d with nulls for %s'
                          % (part_number, self._original_size, end, self.path))
            original_buf = self._download(self.path, start, self._original_size - start)
            modified_buf = bytearray(end - start)
            modified_buf[0:len(original_buf)] = original_buf
            self._mpu.upload_part(modified_buf, offset, part_number)
        else:
            self._mpu.upload_copy_part(start, end, offset, part_number)


class DownloadingMultipartCopyUpload(MultipartUploadDecorator):

    def __init__(self, mpu, download):
        """
        Downloading multipart copy upload.

        Downloads all copy parts from original file and uploads them as regular parts.

        :param mpu: Wrapping multipart upload.
        :param download: Function that retrieves content from an object by its path, offset and length.
        """
        super(DownloadingMultipartCopyUpload, self).__init__(mpu)
        self._mpu = mpu
        self._download = download

    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        self._mpu.upload_part(self._download(self.path, start, end - start), offset, part_number, part_path, keep)


class AppendOptimizedCompositeMultipartCopyUpload(MultipartUploadDecorator):

    def __init__(self, mpu, original_size, chunk_size, download):
        """
        Append optimized composite multipart copy upload.

        Uses original object as a pre uploaded composite part in case of append writes.
        In order to do so it adjusts the first of the already uploaded chunks.

        Uploads copy parts as regular parts using content of the original file.

        :param mpu: Wrapping composite multipart upload.
        :param original_size: Destination file original size.
        :param chunk_size: Size of a single upload part.
        :param download: Function that retrieves content from an object by its path, offset and length.
        """
        super(AppendOptimizedCompositeMultipartCopyUpload, self).__init__(mpu)
        self._mpu = mpu
        self._original_size = original_size
        self._chunk_size = chunk_size
        self._download = download
        self._copy_parts = []
        self._first_chunk = sys.maxint
        self._first_chunk_offset = 0

    def upload_part(self, buf, offset=None, part_number=None, part_path=None, keep=False):
        self._mpu.upload_part(buf, offset, part_number, part_path)
        if part_number < self._first_chunk:
            self._first_chunk = part_number
            self._first_chunk_offset = offset

    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        self._copy_parts.append(_CopyPart(start, end, offset, part_number, part_path, keep))

    def complete(self):
        prefix_copy_part = self._extract_prefix()
        if prefix_copy_part:
            first_uploaded_part_adjusted = self._adjust_first_uploaded_part(self._copy_parts[-1].end)
            if first_uploaded_part_adjusted:
                self._upload_copy_part(prefix_copy_part)
            else:
                self._upload_parts(self._copy_parts)
        else:
            self._upload_parts(self._copy_parts)
        self._mpu.complete()

    def _extract_prefix(self):
        if self._copy_parts:
            sorted_copy_parts = sorted(self._copy_parts, key=lambda copy_part: copy_part.offset)
            prefix_end = 0
            for copy_index, copy_part in enumerate(sorted_copy_parts):
                if copy_part.offset != prefix_end:
                    break
                prefix_end += copy_part.end - copy_part.start
            if prefix_end and prefix_end + self._chunk_size >= self._original_size:
                return _CopyPart(0, self._original_size, offset=0, part_number=1, part_path=self.path, keep=True)
        return None

    def _adjust_first_uploaded_part(self, prefix_end):
        diff_offset = self._original_size - prefix_end
        if diff_offset > 0:
            part_path = self._mpu.composite_part_path(self._first_chunk)
            original_part = self._download(self.path, self._first_chunk_offset, self._chunk_size)
            uploaded_part = self._download(part_path, 0, self._chunk_size)
            if original_part[:diff_offset] == uploaded_part[:diff_offset]:
                self._mpu.upload_part(uploaded_part[diff_offset:],
                                      offset=self._first_chunk_offset, part_number=self._first_chunk,
                                      part_path=part_path)
            else:
                return False
        return True

    def _upload_copy_part(self, copy_part):
        self._mpu.upload_copy_part(copy_part.start, copy_part.end, offset=copy_part.offset,
                                   part_number=copy_part.part_number, part_path=copy_part.part_path,
                                   keep=copy_part.keep)

    def _upload_parts(self, copy_parts):
        for copy_part in copy_parts:
            self._mpu.upload_part(self._download(self.path, copy_part.start, copy_part.end - copy_part.start),
                                  offset=copy_part.offset, part_number=copy_part.part_number,
                                  part_path=copy_part.part_path, keep=copy_part.keep)


class CompositeMultipartUpload(MultipartUpload):

    def __init__(self, bucket, path, new_mpu, mv, max_composite_parts=32):
        """
        Composite object multipart upload.
        
        Uploads each part as an independent file then merges all the files into a single composite one.
        
        Takes into the account maximum allowed composite object parts.

        :param bucket: Destination bucket name.
        :param path: Destination bucket relative path.
        :param new_mpu: Function that instantiates new multipart upload.
        :param mv: Function that moves object within bucket.
        :param max_composite_parts: Number of allowed composite object parts.
        """
        self._bucket = bucket
        self._path = path
        self._new_mpu = new_mpu
        self._mv = mv
        self._max_composite_parts = max_composite_parts
        self._bucket_object = None
        self._blob_object = None
        self._mpus = {}

    @property
    def path(self):
        return self._path

    def initiate(self):
        pass

    def upload_part(self, buf, offset=None, part_number=None, part_path=None, keep=False):
        mpu_number = self._mpu_number(part_number)
        mpu = self._get_mpu(mpu_number)
        mpu.upload_part(buf, offset, part_number, self._part_path(mpu_number, part_number), keep)

    def _mpu_number(self, part_number):
        return (part_number - 1) / self._max_composite_parts + 1

    def _get_mpu(self, mpu_number):
        mpu = self._mpus.get(mpu_number, None)
        if not mpu:
            mpu_path = self._mpu_path(mpu_number)
            mpu = self._new_mpu(mpu_path)
            mpu.initiate()
            self._mpus[mpu_number] = mpu
        return mpu

    def upload_copy_part(self, start, end, offset=None, part_number=None, part_path=None, keep=False):
        mpu_number = self._mpu_number(part_number)
        mpu = self._get_mpu(mpu_number)
        mpu.upload_copy_part(start, end, offset, part_number, part_path, keep)

    def complete(self):
        for mpu in self._mpus.values():
            mpu.complete()
        self._compose([(mpu_number, self._mpus[mpu_number]) for mpu_number in sorted(self._mpus.keys())])

    def _compose(self, mpus):
        if not mpus:
            return
        remaining_mpus = list(mpus)
        composed_mpu = None
        while remaining_mpus:
            if composed_mpu:
                merging_mpus = remaining_mpus[:self._max_composite_parts - 1]
                composed_mpu = self._merge([(0, composed_mpu)] + merging_mpus)
                remaining_mpus = remaining_mpus[self._max_composite_parts - 1:]
            else:
                merging_mpus = remaining_mpus[:self._max_composite_parts]
                composed_mpu = self._merge(merging_mpus)
                remaining_mpus = remaining_mpus[self._max_composite_parts:]
            for mpu_number, _ in merging_mpus:
                del self._mpus[mpu_number]
        self._mv(composed_mpu.path, self._path)

    def _merge(self, mpus):
        merged_mpu = self._new_mpu(self._merged_mpu_path(mpus[0][0], mpus[-1][0]))
        merged_mpu.initiate()
        for mpu_number, mpu in mpus:
            merged_mpu.upload_copy_part(None, None, None, mpu_number, mpu.path)
        merged_mpu.complete()
        return merged_mpu

    def abort(self):
        for mpu in self._mpus.values():
            mpu.abort()

    def composite_part_path(self, part_number):
        return self._part_path(self._mpu_number(part_number), part_number)

    def _mpu_path(self, mpu_number):
        return '%s_%s.tmp/%d' % (self._path, self._hashed_path(), mpu_number)

    def _part_path(self, mpu_number, part_number):
        return '%s_%s.tmp/%d.%d' % (self._path, self._hashed_path(), mpu_number, part_number)

    def _merged_mpu_path(self, left_mpu_number, right_mpu_number):
        return '%s_%s.tmp/%d:%d' % (self._path, self._hashed_path(), left_mpu_number, right_mpu_number)

    def _hashed_path(self):
        return str(abs(hash(self._path)))


class TruncatingMultipartCopyUpload(MultipartUploadDecorator):

    def __init__(self, mpu, length, min_part_number=1):
        """
        Truncating multipart copy upload.

        Truncates the file length to the given size which has to be smaller then the original one.

        :param mpu: Wrapping multipart upload.
        :param length: Target size of the truncating file.
        :param min_part_number: Minimal allowed part number. Same as minimum allowed chunk number.
        """
        super(TruncatingMultipartCopyUpload, self).__init__(mpu)
        self._mpu = mpu
        self._length = length
        self._min_part_number = min_part_number

    def complete(self):
        self._mpu.upload_copy_part(start=0, end=self._length, offset=0, part_number=self._min_part_number)
        self._mpu.complete()