1# -*- coding: utf-8 -*-
2# Copyright 2014 Google Inc. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15"""Helper functions for progress callbacks."""
16
17import logging
18import sys
19
20from gslib.util import MakeHumanReadable
21from gslib.util import UTF8
22
23# Default upper and lower bounds for progress callback frequency.
24_START_BYTES_PER_CALLBACK = 1024*64
25_MAX_BYTES_PER_CALLBACK = 1024*1024*100
26
27# Max width of URL to display in progress indicator. Wide enough to allow
28# 15 chars for x/y display on an 80 char wide terminal.
29MAX_PROGRESS_INDICATOR_COLUMNS = 65
30
31
32class ProgressCallbackWithBackoff(object):
33  """Makes progress callbacks with exponential backoff to a maximum value.
34
35  This prevents excessive log message output.
36  """
37
38  def __init__(self, total_size, callback_func,
39               start_bytes_per_callback=_START_BYTES_PER_CALLBACK,
40               max_bytes_per_callback=_MAX_BYTES_PER_CALLBACK,
41               calls_per_exponent=10):
42    """Initializes the callback with backoff.
43
44    Args:
45      total_size: Total bytes to process. If this is None, size is not known
46          at the outset.
47      callback_func: Func of (int: processed_so_far, int: total_bytes)
48          used to make callbacks.
49      start_bytes_per_callback: Lower bound of bytes per callback.
50      max_bytes_per_callback: Upper bound of bytes per callback.
51      calls_per_exponent: Number of calls to make before reducing rate.
52    """
53    self._bytes_per_callback = start_bytes_per_callback
54    self._callback_func = callback_func
55    self._calls_per_exponent = calls_per_exponent
56    self._max_bytes_per_callback = max_bytes_per_callback
57    self._total_size = total_size
58
59    self._bytes_processed_since_callback = 0
60    self._callbacks_made = 0
61    self._total_bytes_processed = 0
62
63  def Progress(self, bytes_processed):
64    """Tracks byte processing progress, making a callback if necessary."""
65    self._bytes_processed_since_callback += bytes_processed
66    if (self._bytes_processed_since_callback > self._bytes_per_callback or
67        (self._total_bytes_processed + self._bytes_processed_since_callback >=
68         self._total_size and self._total_size is not None)):
69      self._total_bytes_processed += self._bytes_processed_since_callback
70      # TODO: We check if >= total_size and truncate because JSON uploads count
71      # headers+metadata during their send progress. If the size is unknown,
72      # we can't do this and the progress message will make it appear that we
73      # send more than the original stream.
74      if self._total_size is not None:
75        bytes_sent = min(self._total_bytes_processed, self._total_size)
76      else:
77        bytes_sent = self._total_bytes_processed
78      self._callback_func(bytes_sent, self._total_size)
79      self._bytes_processed_since_callback = 0
80      self._callbacks_made += 1
81
82      if self._callbacks_made > self._calls_per_exponent:
83        self._bytes_per_callback = min(self._bytes_per_callback * 2,
84                                       self._max_bytes_per_callback)
85        self._callbacks_made = 0
86
87
88def ConstructAnnounceText(operation_name, url_string):
89  """Constructs announce text for ongoing operations on url_to_display.
90
91  This truncates the text to a maximum of MAX_PROGRESS_INDICATOR_COLUMNS.
92  Thus, concurrent output (gsutil -m) leaves progress counters in a readable
93  (fixed) position.
94
95  Args:
96    operation_name: String describing the operation, i.e.
97        'Uploading' or 'Hashing'.
98    url_string: String describing the file/object being processed.
99
100  Returns:
101    Formatted announce text for outputting operation progress.
102  """
103  # Operation name occupies 11 characters (enough for 'Downloading'), plus a
104  # space. The rest is used for url_to_display. If a longer operation name is
105  # used, it will be truncated. We can revisit this size if we need to support
106  # a longer operation, but want to make sure the terminal output is meaningful.
107  justified_op_string = operation_name[:11].ljust(12)
108  start_len = len(justified_op_string)
109  end_len = len(': ')
110  if (start_len + len(url_string) + end_len >
111      MAX_PROGRESS_INDICATOR_COLUMNS):
112    ellipsis_len = len('...')
113    url_string = '...%s' % url_string[
114        -(MAX_PROGRESS_INDICATOR_COLUMNS - start_len - end_len - ellipsis_len):]
115  base_announce_text = '%s%s:' % (justified_op_string, url_string)
116  format_str = '{0:%ds}' % MAX_PROGRESS_INDICATOR_COLUMNS
117  return format_str.format(base_announce_text.encode(UTF8))
118
119
120class FileProgressCallbackHandler(object):
121  """Outputs progress info for large operations like file copy or hash."""
122
123  def __init__(self, announce_text, logger, start_byte=0,
124               override_total_size=None):
125    """Initializes the callback handler.
126
127    Args:
128      announce_text: String describing the operation.
129      logger: For outputting log messages.
130      start_byte: The beginning of the file component, if one is being used.
131      override_total_size: The size of the file component, if one is being used.
132    """
133    self._announce_text = announce_text
134    self._logger = logger
135    self._start_byte = start_byte
136    self._override_total_size = override_total_size
137    # Ensures final newline is written once even if we get multiple callbacks.
138    self._last_byte_written = False
139
140  # Function signature is in boto callback format, which cannot be changed.
141  def call(self,  # pylint: disable=invalid-name
142           last_byte_processed,
143           total_size):
144    """Prints an overwriting line to stderr describing the operation progress.
145
146    Args:
147      last_byte_processed: The last byte processed in the file. For file
148                           components, this number should be in the range
149                           [start_byte:start_byte + override_total_size].
150      total_size: Total size of the ongoing operation.
151    """
152    if not self._logger.isEnabledFor(logging.INFO) or self._last_byte_written:
153      return
154
155    if self._override_total_size:
156      total_size = self._override_total_size
157
158    if total_size:
159      total_size_string = '/%s' % MakeHumanReadable(total_size)
160    else:
161      total_size_string = ''
162    # Use sys.stderr.write instead of self.logger.info so progress messages
163    # output on a single continuously overwriting line.
164    # TODO: Make this work with logging.Logger.
165    sys.stderr.write('%s%s%s    \r' % (
166        self._announce_text,
167        MakeHumanReadable(last_byte_processed - self._start_byte),
168        total_size_string))
169    if total_size and last_byte_processed - self._start_byte == total_size:
170      self._last_byte_written = True
171      sys.stderr.write('\n')
172