Fix deduplication report test.
[arvados.git] / lib / deduplicationreport / report.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package deduplicationreport
6
7 import (
8         "flag"
9         "fmt"
10         "io"
11         "strings"
12
13         "git.arvados.org/arvados.git/sdk/go/arvados"
14         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
15         "git.arvados.org/arvados.git/sdk/go/manifest"
16
17         "github.com/dustin/go-humanize"
18         "github.com/sirupsen/logrus"
19 )
20
21 func deDuplicate(inputs []string) (trimmed []string) {
22         seen := make(map[string]bool)
23         for _, uuid := range inputs {
24                 if !seen[uuid] {
25                         seen[uuid] = true
26                         trimmed = append(trimmed, uuid)
27                 }
28         }
29         return
30 }
31
32 func parseFlags(prog string, args []string, logger *logrus.Logger, stderr io.Writer) ([]string, error) {
33         flags := flag.NewFlagSet("", flag.ContinueOnError)
34         flags.SetOutput(stderr)
35         flags.Usage = func() {
36                 fmt.Fprintf(flags.Output(), `
37 Usage:
38   %s [options ...] <collection-uuid> <collection-uuid> ...
39
40   %s [options ...] <collection-pdh>,<collection_uuid> \
41      <collection-pdh>,<collection_uuid> ...
42
43   This program analyzes the overlap in blocks used by 2 or more collections. It
44   prints a deduplication report that shows the nominal space used by the
45   collections, as well as the actual size and the amount of space that is saved
46   by Keep's deduplication.
47
48   The list of collections may be provided in two ways. A list of collection
49   uuids is sufficient. Alternatively, the PDH for each collection may also be
50   provided. This is will greatly speed up operation when the list contains
51   multiple collections with the same PDH.
52
53   Exit status will be zero if there were no errors generating the report.
54
55 Example:
56
57   Use the 'arv' and 'jq' commands to get the list of the 100
58   largest collections and generate the deduplication report:
59
60   arv collection list --order 'file_size_total desc' --limit 100 | \
61     jq -r '.items[] | [.portable_data_hash,.uuid] |@csv' | \
62     sed -e 's/"//g'|tr '\n' ' ' | \
63     xargs %s
64
65 Options:
66 `, prog, prog, prog)
67                 flags.PrintDefaults()
68         }
69         loglevel := flags.String("log-level", "info", "logging level (debug, info, ...)")
70         err := flags.Parse(args)
71         if err == flag.ErrHelp {
72                 return nil, err
73         } else if err != nil {
74                 return nil, err
75         }
76
77         inputs := flags.Args()
78
79         inputs = deDuplicate(inputs)
80
81         if len(inputs) < 1 {
82                 err = fmt.Errorf("Error: no collections provided")
83                 return inputs, err
84         }
85
86         lvl, err := logrus.ParseLevel(*loglevel)
87         if err != nil {
88                 return inputs, err
89         }
90         logger.SetLevel(lvl)
91         return inputs, err
92 }
93
94 func blockList(collection arvados.Collection) (blocks map[string]int) {
95         blocks = make(map[string]int)
96         m := manifest.Manifest{Text: collection.ManifestText}
97         blockChannel := m.BlockIterWithDuplicates()
98         for b := range blockChannel {
99                 blocks[b.Digest.String()] = b.Size
100         }
101         return
102 }
103
104 func report(prog string, args []string, logger *logrus.Logger, stdout, stderr io.Writer) (exitcode int) {
105         var inputs []string
106         var err error
107
108         inputs, err = parseFlags(prog, args, logger, stderr)
109         if err != nil {
110                 if err != flag.ErrHelp {
111                         logger.Error(err.Error())
112                         return 2
113                 }
114                 return 0
115         }
116
117         // Arvados Client setup
118         arv, err := arvadosclient.MakeArvadosClient()
119         if err != nil {
120                 logger.Errorf("Error creating Arvados object: %s", err)
121                 exitcode = 1
122                 return
123         }
124
125         type Col struct {
126                 FileSizeTotal int64
127                 FileCount     int64
128         }
129
130         blocks := make(map[string]map[string]int)
131         pdhs := make(map[string]Col)
132         var nominalSize int64
133
134         for _, input := range inputs {
135                 var uuid string
136                 var pdh string
137                 if strings.Contains(input, ",") {
138                         // The input is in the format pdh,uuid. This will allow us to save time on duplicate pdh's
139                         tmp := strings.Split(input, ",")
140                         pdh = tmp[0]
141                         uuid = tmp[1]
142                 } else {
143                         // The input must be a plain uuid
144                         uuid = input
145                 }
146                 if !strings.Contains(uuid, "-4zz18-") {
147                         logger.Errorf("Error: uuid must refer to collection object")
148                         exitcode = 1
149                         return
150                 }
151                 if _, ok := pdhs[pdh]; ok {
152                         // We've processed a collection with this pdh already. Simply add its
153                         // size to the totals and move on to the next one.
154                         // Note that we simply trust the PDH matches the collection UUID here,
155                         // in other words, we use it over the UUID. If they don't match, the report
156                         // will be wrong.
157                         nominalSize += pdhs[pdh].FileSizeTotal
158                 } else {
159                         var collection arvados.Collection
160                         err = arv.Get("collections", uuid, nil, &collection)
161                         if err != nil {
162                                 logger.Errorf("Error: unable to retrieve collection: %s", err)
163                                 exitcode = 1
164                                 return
165                         }
166                         blocks[uuid] = make(map[string]int)
167                         blocks[uuid] = blockList(collection)
168                         if pdh != "" && collection.PortableDataHash != pdh {
169                                 logger.Errorf("Error: the collection with UUID %s has PDH %s, but a different PDH was provided in the arguments: %s", uuid, collection.PortableDataHash, pdh)
170                                 exitcode = 1
171                                 return
172                         }
173                         if pdh == "" {
174                                 pdh = collection.PortableDataHash
175                         }
176
177                         col := Col{}
178                         if collection.FileSizeTotal != 0 || collection.FileCount != 0 {
179                                 nominalSize += collection.FileSizeTotal
180                                 col.FileSizeTotal = collection.FileSizeTotal
181                                 col.FileCount = int64(collection.FileCount)
182                         } else {
183                                 // Collections created with old Arvados versions do not always have the total file size and count cached in the collections object
184                                 var collSize int64
185                                 for _, size := range blocks[uuid] {
186                                         collSize += int64(size)
187                                 }
188                                 nominalSize += collSize
189                                 col.FileSizeTotal = collSize
190                         }
191                         pdhs[pdh] = col
192                 }
193
194                 if pdhs[pdh].FileCount != 0 {
195                         fmt.Fprintf(stdout, "Collection %s: pdh %s; nominal size %d (%s); file count %d\n", uuid, pdh, pdhs[pdh].FileSizeTotal, humanize.IBytes(uint64(pdhs[pdh].FileSizeTotal)), pdhs[pdh].FileCount)
196                 } else {
197                         fmt.Fprintf(stdout, "Collection %s: pdh %s; nominal size %d (%s)\n", uuid, pdh, pdhs[pdh].FileSizeTotal, humanize.IBytes(uint64(pdhs[pdh].FileSizeTotal)))
198                 }
199         }
200
201         var totalSize int64
202         seen := make(map[string]bool)
203         for _, v := range blocks {
204                 for pdh, size := range v {
205                         if !seen[pdh] {
206                                 seen[pdh] = true
207                                 totalSize += int64(size)
208                         }
209                 }
210         }
211         fmt.Fprintln(stdout)
212         fmt.Fprintf(stdout, "Collections:                 %15d\n", len(inputs))
213         fmt.Fprintf(stdout, "Nominal size of stored data: %15d bytes (%s)\n", nominalSize, humanize.IBytes(uint64(nominalSize)))
214         fmt.Fprintf(stdout, "Actual size of stored data:  %15d bytes (%s)\n", totalSize, humanize.IBytes(uint64(totalSize)))
215         fmt.Fprintf(stdout, "Saved by Keep deduplication: %15d bytes (%s)\n", nominalSize-totalSize, humanize.IBytes(uint64(nominalSize-totalSize)))
216
217         return exitcode
218 }