Modifies sanity check to query for the total expected number of collections
[arvados.git] / services / datamanager / collection / collection.go
1 // Deals with parsing Collection responses from API Server.
2
3 package collection
4
5 import (
6         "flag"
7         "fmt"
8         "git.curoverse.com/arvados.git/sdk/go/arvadosclient"
9         "git.curoverse.com/arvados.git/sdk/go/blockdigest"
10         "git.curoverse.com/arvados.git/sdk/go/logger"
11         "git.curoverse.com/arvados.git/sdk/go/manifest"
12         "git.curoverse.com/arvados.git/sdk/go/util"
13         "log"
14         "os"
15         "runtime/pprof"
16         "time"
17 )
18
19 var (
20         HeapProfileFilename string
21 )
22
23 // Collection representation
24 type Collection struct {
25         UUID              string
26         OwnerUUID         string
27         ReplicationLevel  int
28         BlockDigestToSize map[blockdigest.BlockDigest]int
29         TotalSize         int
30 }
31
32 // ReadCollections holds information about collections from API server
33 type ReadCollections struct {
34         ReadAllCollections        bool
35         UUIDToCollection          map[string]Collection
36         OwnerToCollectionSize     map[string]int
37         BlockToDesiredReplication map[blockdigest.DigestWithSize]int
38         CollectionUUIDToIndex     map[string]int
39         CollectionIndexToUUID     []string
40         BlockToCollectionIndices  map[blockdigest.DigestWithSize][]int
41 }
42
43 // GetCollectionsParams params
44 type GetCollectionsParams struct {
45         Client    arvadosclient.ArvadosClient
46         Logger    *logger.Logger
47         BatchSize int
48 }
49
50 // SdkCollectionInfo holds collection info from api
51 type SdkCollectionInfo struct {
52         UUID         string    `json:"uuid"`
53         OwnerUUID    string    `json:"owner_uuid"`
54         Redundancy   int       `json:"redundancy"`
55         ModifiedAt   time.Time `json:"modified_at"`
56         ManifestText string    `json:"manifest_text"`
57 }
58
59 // SdkCollectionList lists collections from api
60 type SdkCollectionList struct {
61         ItemsAvailable int                 `json:"items_available"`
62         Items          []SdkCollectionInfo `json:"items"`
63 }
64
65 func init() {
66         flag.StringVar(&HeapProfileFilename,
67                 "heap-profile",
68                 "",
69                 "File to write the heap profiles to. Leave blank to skip profiling.")
70 }
71
72 // WriteHeapProfile writes the heap profile to a file for later review.
73 // Since a file is expected to only contain a single heap profile this
74 // function overwrites the previously written profile, so it is safe
75 // to call multiple times in a single run.
76 // Otherwise we would see cumulative numbers as explained here:
77 // https://groups.google.com/d/msg/golang-nuts/ZyHciRglQYc/2nh4Ndu2fZcJ
78 func WriteHeapProfile() error {
79         if HeapProfileFilename != "" {
80                 heapProfile, err := os.Create(HeapProfileFilename)
81                 if err != nil {
82                         return err
83                 }
84
85                 defer heapProfile.Close()
86
87                 err = pprof.WriteHeapProfile(heapProfile)
88                 return err
89         }
90
91         return nil
92 }
93
94 // GetCollectionsAndSummarize gets collections from api and summarizes
95 func GetCollectionsAndSummarize(params GetCollectionsParams) (results ReadCollections, err error) {
96         results, err = GetCollections(params)
97         if err != nil {
98                 return
99         }
100
101         results.Summarize(params.Logger)
102
103         log.Printf("Uuid to Size used: %v", results.OwnerToCollectionSize)
104         log.Printf("Read and processed %d collections",
105                 len(results.UUIDToCollection))
106
107         // TODO(misha): Add a "readonly" flag. If we're in readonly mode,
108         // lots of behaviors can become warnings (and obviously we can't
109         // write anything).
110         // if !readCollections.ReadAllCollections {
111         //      log.Fatalf("Did not read all collections")
112         // }
113
114         return
115 }
116
117 // GetCollections gets collections from api
118 func GetCollections(params GetCollectionsParams) (results ReadCollections, err error) {
119         if &params.Client == nil {
120                 err = fmt.Errorf("params.Client passed to GetCollections() should " +
121                         "contain a valid ArvadosClient, but instead it is nil.")
122                 return
123         }
124
125         fieldsWanted := []string{"manifest_text",
126                 "owner_uuid",
127                 "uuid",
128                 "redundancy",
129                 "modified_at"}
130
131         sdkParams := arvadosclient.Dict{
132                 "select":  fieldsWanted,
133                 "order":   []string{"modified_at ASC"},
134                 "filters": [][]string{[]string{"modified_at", ">=", "1900-01-01T00:00:00Z"}}}
135
136         if params.BatchSize > 0 {
137                 sdkParams["limit"] = params.BatchSize
138         }
139
140         var defaultReplicationLevel int
141         {
142                 var value interface{}
143                 value, err = params.Client.Discovery("defaultCollectionReplication")
144                 if err != nil {
145                         return
146                 }
147
148                 defaultReplicationLevel = int(value.(float64))
149                 if defaultReplicationLevel <= 0 {
150                         err = fmt.Errorf("Default collection replication returned by arvados SDK "+
151                                 "should be a positive integer but instead it was %d.",
152                                 defaultReplicationLevel)
153                         return
154                 }
155         }
156
157         initialNumberOfCollectionsAvailable, err :=
158                 util.NumberItemsAvailable(params.Client, "collections")
159         if err != nil {
160                 return
161         }
162         // Include a 1% margin for collections added while we're reading so
163         // that we don't have to grow the map in most cases.
164         maxExpectedCollections := int(
165                 float64(initialNumberOfCollectionsAvailable) * 1.01)
166         results.UUIDToCollection = make(map[string]Collection, maxExpectedCollections)
167
168         if params.Logger != nil {
169                 params.Logger.Update(func(p map[string]interface{}, e map[string]interface{}) {
170                         collectionInfo := logger.GetOrCreateMap(p, "collection_info")
171                         collectionInfo["num_collections_at_start"] = initialNumberOfCollectionsAvailable
172                         collectionInfo["batch_size"] = params.BatchSize
173                         collectionInfo["default_replication_level"] = defaultReplicationLevel
174                 })
175         }
176
177         // These values are just for getting the loop to run the first time,
178         // afterwards they'll be set to real values.
179         previousTotalCollections := -1
180         totalCollections := 0
181         for totalCollections > previousTotalCollections {
182                 // We're still finding new collections
183
184                 // Write the heap profile for examining memory usage
185                 err = WriteHeapProfile()
186                 if err != nil {
187                         return
188                 }
189
190                 // Get next batch of collections.
191                 var collections SdkCollectionList
192                 err = params.Client.List("collections", sdkParams, &collections)
193                 if err != nil {
194                         return
195                 }
196
197                 // Process collection and update our date filter.
198                 latestModificationDate, maxManifestSize, totalManifestSize, err := ProcessCollections(params.Logger,
199                         collections.Items,
200                         defaultReplicationLevel,
201                         results.UUIDToCollection)
202                 if err != nil {
203                         return results, err
204                 }
205                 sdkParams["filters"].([][]string)[0][2] = latestModificationDate.Format(time.RFC3339)
206
207                 // update counts
208                 previousTotalCollections = totalCollections
209                 totalCollections = len(results.UUIDToCollection)
210
211                 log.Printf("%d collections read, %d new in last batch, "+
212                         "%s latest modified date, %.0f %d %d avg,max,total manifest size",
213                         totalCollections,
214                         totalCollections-previousTotalCollections,
215                         sdkParams["filters"].([][]string)[0][2],
216                         float32(totalManifestSize)/float32(totalCollections),
217                         maxManifestSize, totalManifestSize)
218
219                 if params.Logger != nil {
220                         params.Logger.Update(func(p map[string]interface{}, e map[string]interface{}) {
221                                 collectionInfo := logger.GetOrCreateMap(p, "collection_info")
222                                 collectionInfo["collections_read"] = totalCollections
223                                 collectionInfo["latest_modified_date_seen"] = sdkParams["filters"].([][]string)[0][2]
224                                 collectionInfo["total_manifest_size"] = totalManifestSize
225                                 collectionInfo["max_manifest_size"] = maxManifestSize
226                         })
227                 }
228         }
229
230         // Make one final API request to verify that we have processed all collections available up to the latest modification date
231         sdkParams["filters"].([][]string)[0][1] = "<="
232         sdkParams["limit"] = 0
233         err = params.Client.List("collections", sdkParams, &collections)
234         if err != nil {
235                 return
236         }
237         finalNumberOfCollectionsAvailable, err :=
238                 util.NumberItemsAvailable(params.Client, "collections")
239         if err != nil {
240                 return
241         }
242         if totalCollections < finalNumberOfCollectionsAvailable {
243                 err = fmt.Errorf("API server indicates a total of %d collections "+
244                                 "available up to %v, but we only retrieved %d. "+
245                                 "Refusing to continue as this could indicate an "+
246                                 "otherwise undetected failure.",
247                                 finalNumberOfCollectionsAvailable, 
248                                 sdkParams["filters"].([][]string)[0][2],
249                                 totalCollections)
250                 return
251         }
252
253         // Write the heap profile for examining memory usage
254         err = WriteHeapProfile()
255
256         return
257 }
258
259 // StrCopy returns a newly allocated string.
260 // It is useful to copy slices so that the garbage collector can reuse
261 // the memory of the longer strings they came from.
262 func StrCopy(s string) string {
263         return string([]byte(s))
264 }
265
266 // ProcessCollections read from api server
267 func ProcessCollections(arvLogger *logger.Logger,
268         receivedCollections []SdkCollectionInfo,
269         defaultReplicationLevel int,
270         UUIDToCollection map[string]Collection,
271 ) (
272         latestModificationDate time.Time,
273         maxManifestSize, totalManifestSize uint64,
274         err error,
275 ) {
276         for _, sdkCollection := range receivedCollections {
277                 collection := Collection{UUID: StrCopy(sdkCollection.UUID),
278                         OwnerUUID:         StrCopy(sdkCollection.OwnerUUID),
279                         ReplicationLevel:  sdkCollection.Redundancy,
280                         BlockDigestToSize: make(map[blockdigest.BlockDigest]int)}
281
282                 if sdkCollection.ModifiedAt.IsZero() {
283                         err = fmt.Errorf(
284                                 "Arvados SDK collection returned with unexpected zero "+
285                                         "modification date. This probably means that either we failed to "+
286                                         "parse the modification date or the API server has changed how "+
287                                         "it returns modification dates: %+v",
288                                 collection)
289                         return
290                 }
291
292                 if sdkCollection.ModifiedAt.After(latestModificationDate) {
293                         latestModificationDate = sdkCollection.ModifiedAt
294                 }
295
296                 if collection.ReplicationLevel == 0 {
297                         collection.ReplicationLevel = defaultReplicationLevel
298                 }
299
300                 manifest := manifest.Manifest{Text: sdkCollection.ManifestText}
301                 manifestSize := uint64(len(sdkCollection.ManifestText))
302
303                 if _, alreadySeen := UUIDToCollection[collection.UUID]; !alreadySeen {
304                         totalManifestSize += manifestSize
305                 }
306                 if manifestSize > maxManifestSize {
307                         maxManifestSize = manifestSize
308                 }
309
310                 blockChannel := manifest.BlockIterWithDuplicates()
311                 for block := range blockChannel {
312                         if storedSize, stored := collection.BlockDigestToSize[block.Digest]; stored && storedSize != block.Size {
313                                 log.Printf(
314                                         "Collection %s contains multiple sizes (%d and %d) for block %s",
315                                         collection.UUID,
316                                         storedSize,
317                                         block.Size,
318                                         block.Digest)
319                         }
320                         collection.BlockDigestToSize[block.Digest] = block.Size
321                 }
322                 if manifest.Err != nil {
323                         err = manifest.Err
324                         return
325                 }
326
327                 collection.TotalSize = 0
328                 for _, size := range collection.BlockDigestToSize {
329                         collection.TotalSize += size
330                 }
331                 UUIDToCollection[collection.UUID] = collection
332
333                 // Clear out all the manifest strings that we don't need anymore.
334                 // These hopefully form the bulk of our memory usage.
335                 manifest.Text = ""
336                 sdkCollection.ManifestText = ""
337         }
338
339         return
340 }
341
342 // Summarize the collections read
343 func (readCollections *ReadCollections) Summarize(arvLogger *logger.Logger) {
344         readCollections.OwnerToCollectionSize = make(map[string]int)
345         readCollections.BlockToDesiredReplication = make(map[blockdigest.DigestWithSize]int)
346         numCollections := len(readCollections.UUIDToCollection)
347         readCollections.CollectionUUIDToIndex = make(map[string]int, numCollections)
348         readCollections.CollectionIndexToUUID = make([]string, 0, numCollections)
349         readCollections.BlockToCollectionIndices = make(map[blockdigest.DigestWithSize][]int)
350
351         for _, coll := range readCollections.UUIDToCollection {
352                 collectionIndex := len(readCollections.CollectionIndexToUUID)
353                 readCollections.CollectionIndexToUUID =
354                         append(readCollections.CollectionIndexToUUID, coll.UUID)
355                 readCollections.CollectionUUIDToIndex[coll.UUID] = collectionIndex
356
357                 readCollections.OwnerToCollectionSize[coll.OwnerUUID] =
358                         readCollections.OwnerToCollectionSize[coll.OwnerUUID] + coll.TotalSize
359
360                 for block, size := range coll.BlockDigestToSize {
361                         locator := blockdigest.DigestWithSize{Digest: block, Size: uint32(size)}
362                         readCollections.BlockToCollectionIndices[locator] =
363                                 append(readCollections.BlockToCollectionIndices[locator],
364                                         collectionIndex)
365                         storedReplication := readCollections.BlockToDesiredReplication[locator]
366                         if coll.ReplicationLevel > storedReplication {
367                                 readCollections.BlockToDesiredReplication[locator] =
368                                         coll.ReplicationLevel
369                         }
370                 }
371         }
372
373         if arvLogger != nil {
374                 arvLogger.Update(func(p map[string]interface{}, e map[string]interface{}) {
375                         collectionInfo := logger.GetOrCreateMap(p, "collection_info")
376                         // Since maps are shallow copied, we run a risk of concurrent
377                         // updates here. By copying results.OwnerToCollectionSize into
378                         // the log, we're assuming that it won't be updated.
379                         collectionInfo["owner_to_collection_size"] =
380                                 readCollections.OwnerToCollectionSize
381                         collectionInfo["distinct_blocks_named"] =
382                                 len(readCollections.BlockToDesiredReplication)
383                 })
384         }
385
386         return
387 }