Merge branch '19316-oj-safe-load'
[arvados.git] / sdk / python / arvados / _ranges.py
1 # Copyright (C) The Arvados Authors. All rights reserved.
2 #
3 # SPDX-License-Identifier: Apache-2.0
4
5 from __future__ import division
6 from builtins import object
7 import logging
8
9 _logger = logging.getLogger('arvados.ranges')
10
11 # Log level below 'debug' !
12 RANGES_SPAM = 9
13
14 class Range(object):
15     __slots__ = ("locator", "range_start", "range_size", "segment_offset")
16
17     def __init__(self, locator, range_start, range_size, segment_offset=0):
18         self.locator = locator
19         self.range_start = range_start
20         self.range_size = range_size
21         self.segment_offset = segment_offset
22
23     def __repr__(self):
24         return "Range(%r, %r, %r, %r)" % (self.locator, self.range_start, self.range_size, self.segment_offset)
25
26     def __eq__(self, other):
27         return (self.locator == other.locator and
28                 self.range_start == other.range_start and
29                 self.range_size == other.range_size and
30                 self.segment_offset == other.segment_offset)
31
32 def first_block(data_locators, range_start):
33     block_start = 0
34
35     # range_start/block_start is the inclusive lower bound
36     # range_end/block_end is the exclusive upper bound
37
38     hi = len(data_locators)
39     lo = 0
40     i = (hi + lo) // 2
41     block_size = data_locators[i].range_size
42     block_start = data_locators[i].range_start
43     block_end = block_start + block_size
44
45     # perform a binary search for the first block
46     # assumes that all of the blocks are contiguous, so range_start is guaranteed
47     # to either fall into the range of a block or be outside the block range entirely
48     while not (range_start >= block_start and range_start < block_end):
49         if lo == i:
50             # must be out of range, fail
51             return None
52         if range_start > block_start:
53             lo = i
54         else:
55             hi = i
56         i = (hi + lo) // 2
57         block_size = data_locators[i].range_size
58         block_start = data_locators[i].range_start
59         block_end = block_start + block_size
60
61     return i
62
63 class LocatorAndRange(object):
64     __slots__ = ("locator", "block_size", "segment_offset", "segment_size")
65
66     def __init__(self, locator, block_size, segment_offset, segment_size):
67         self.locator = locator
68         self.block_size = block_size
69         self.segment_offset = segment_offset
70         self.segment_size = segment_size
71
72     def __eq__(self, other):
73         return  (self.locator == other.locator and
74                  self.block_size == other.block_size and
75                  self.segment_offset == other.segment_offset and
76                  self.segment_size == other.segment_size)
77
78     def __repr__(self):
79         return "LocatorAndRange(%r, %r, %r, %r)" % (self.locator, self.block_size, self.segment_offset, self.segment_size)
80
81 def locators_and_ranges(data_locators, range_start, range_size, limit=None):
82     """Get blocks that are covered by a range.
83
84     Returns a list of LocatorAndRange objects.
85
86     :data_locators:
87       list of Range objects, assumes that blocks are in order and contiguous
88
89     :range_start:
90       start of range
91
92     :range_size:
93       size of range
94
95     :limit:
96       Maximum segments to return, default None (unlimited).  Will truncate the
97       result if there are more segments needed to cover the range than the
98       limit.
99
100     """
101     if range_size == 0:
102         return []
103     resp = []
104     range_end = range_start + range_size
105
106     i = first_block(data_locators, range_start)
107     if i is None:
108         return []
109
110     # We should always start at the first segment due to the binary
111     # search.
112     while i < len(data_locators) and len(resp) != limit:
113         dl = data_locators[i]
114         block_start = dl.range_start
115         block_size = dl.range_size
116         block_end = block_start + block_size
117         _logger.log(RANGES_SPAM,
118             "L&R %s range_start %s block_start %s range_end %s block_end %s",
119             dl.locator, range_start, block_start, range_end, block_end)
120         if range_end <= block_start:
121             # range ends before this block starts, so don't look at any more locators
122             break
123
124         if range_start >= block_start and range_end <= block_end:
125             # range starts and ends in this block
126             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset + (range_start - block_start), range_size))
127         elif range_start >= block_start and range_end > block_end:
128             # range starts in this block
129             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset + (range_start - block_start), block_end - range_start))
130         elif range_start < block_start and range_end > block_end:
131             # range starts in a previous block and extends to further blocks
132             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset, block_size))
133         elif range_start < block_start and range_end <= block_end:
134             # range starts in a previous block and ends in this block
135             resp.append(LocatorAndRange(dl.locator, block_size, dl.segment_offset, range_end - block_start))
136         block_start = block_end
137         i += 1
138     return resp
139
140 def replace_range(data_locators, new_range_start, new_range_size, new_locator, new_segment_offset):
141     """
142     Replace a file segment range with a new segment.
143
144     NOTE::
145     data_locators will be updated in place
146
147     :data_locators:
148       list of Range objects, assumes that segments are in order and contiguous
149
150     :new_range_start:
151       start of range to replace in data_locators
152
153     :new_range_size:
154       size of range to replace in data_locators
155
156     :new_locator:
157       locator for new segment to be inserted
158
159     :new_segment_offset:
160       segment offset within the locator
161
162     """
163     if new_range_size == 0:
164         return
165
166     new_range_end = new_range_start + new_range_size
167
168     if len(data_locators) == 0:
169         data_locators.append(Range(new_locator, new_range_start, new_range_size, new_segment_offset))
170         return
171
172     last = data_locators[-1]
173     if (last.range_start+last.range_size) == new_range_start:
174         if last.locator == new_locator and (last.segment_offset+last.range_size) == new_segment_offset:
175             # extend last segment
176             last.range_size += new_range_size
177         else:
178             data_locators.append(Range(new_locator, new_range_start, new_range_size, new_segment_offset))
179         return
180
181     i = first_block(data_locators, new_range_start)
182     if i is None:
183         return
184
185     # We should always start at the first segment due to the binary
186     # search.
187     while i < len(data_locators):
188         dl = data_locators[i]
189         old_segment_start = dl.range_start
190         old_segment_end = old_segment_start + dl.range_size
191         _logger.log(RANGES_SPAM,
192             "RR %s range_start %s segment_start %s range_end %s segment_end %s",
193             dl, new_range_start, old_segment_start, new_range_end,
194             old_segment_end)
195         if new_range_end <= old_segment_start:
196             # range ends before this segment starts, so don't look at any more locators
197             break
198
199         if old_segment_start <= new_range_start and new_range_end <= old_segment_end:
200             # new range starts and ends in old segment
201             # split segment into up to 3 pieces
202             if (new_range_start-old_segment_start) > 0:
203                 data_locators[i] = Range(dl.locator, old_segment_start, (new_range_start-old_segment_start), dl.segment_offset)
204                 data_locators.insert(i+1, Range(new_locator, new_range_start, new_range_size, new_segment_offset))
205             else:
206                 data_locators[i] = Range(new_locator, new_range_start, new_range_size, new_segment_offset)
207                 i -= 1
208             if (old_segment_end-new_range_end) > 0:
209                 data_locators.insert(i+2, Range(dl.locator, new_range_end, (old_segment_end-new_range_end), dl.segment_offset + (new_range_start-old_segment_start) + new_range_size))
210             return
211         elif old_segment_start <= new_range_start and new_range_end > old_segment_end:
212             # range starts in this segment
213             # split segment into 2 pieces
214             data_locators[i] = Range(dl.locator, old_segment_start, (new_range_start-old_segment_start), dl.segment_offset)
215             data_locators.insert(i+1, Range(new_locator, new_range_start, new_range_size, new_segment_offset))
216             i += 1
217         elif new_range_start < old_segment_start and new_range_end >= old_segment_end:
218             # range starts in a previous segment and extends to further segments
219             # delete this segment
220             del data_locators[i]
221             i -= 1
222         elif new_range_start < old_segment_start and new_range_end < old_segment_end:
223             # range starts in a previous segment and ends in this segment
224             # move the starting point of this segment up, and shrink it.
225             data_locators[i] = Range(dl.locator, new_range_end, (old_segment_end-new_range_end), dl.segment_offset + (new_range_end-old_segment_start))
226             return
227         i += 1