Fix another wb2 integration test typo, refs #9964
[arvados.git] / services / keep-balance / block_state.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package keepbalance
6
7 import (
8         "sync"
9
10         "git.arvados.org/arvados.git/sdk/go/arvados"
11 )
12
13 // Replica is a file on disk (or object in an S3 bucket, or blob in an
14 // Azure storage container, etc.) as reported in a keepstore index
15 // response.
16 type Replica struct {
17         *KeepMount
18         Mtime int64
19 }
20
21 // BlockState indicates the desired storage class and number of
22 // replicas (according to the collections we know about) and the
23 // replicas actually stored (according to the keepstore indexes we
24 // know about).
25 type BlockState struct {
26         Refs     map[string]bool // pdh => true (only tracked when len(Replicas)==0)
27         RefCount int
28         Replicas []Replica
29         Desired  map[string]int
30         // TODO: Support combinations of classes ("private + durable")
31         // by replacing the map[string]int with a map[*[]string]int
32         // here, where the map keys come from a pool of semantically
33         // distinct class combinations.
34         //
35         // TODO: Use a pool of semantically distinct Desired maps to
36         // conserve memory (typically there are far more BlockState
37         // objects in memory than distinct Desired profiles).
38 }
39
40 var defaultClasses = []string{"default"}
41
42 func (bs *BlockState) addReplica(r Replica) {
43         bs.Replicas = append(bs.Replicas, r)
44         // Free up memory wasted by tracking PDHs that will never be
45         // reported (see comment in increaseDesired)
46         bs.Refs = nil
47 }
48
49 func (bs *BlockState) increaseDesired(pdh string, classes []string, n int) {
50         if pdh != "" && len(bs.Replicas) == 0 {
51                 // Note we only track PDHs if there's a possibility
52                 // that we will report the list of referring PDHs,
53                 // i.e., if we haven't yet seen a replica.
54                 if bs.Refs == nil {
55                         bs.Refs = map[string]bool{}
56                 }
57                 bs.Refs[pdh] = true
58         }
59         bs.RefCount++
60         if len(classes) == 0 {
61                 classes = defaultClasses
62         }
63         for _, class := range classes {
64                 if bs.Desired == nil {
65                         bs.Desired = map[string]int{class: n}
66                 } else if d, ok := bs.Desired[class]; !ok || d < n {
67                         bs.Desired[class] = n
68                 }
69         }
70 }
71
72 // BlockStateMap is a goroutine-safe wrapper around a
73 // map[arvados.SizedDigest]*BlockState.
74 type BlockStateMap struct {
75         entries map[arvados.SizedDigest]*BlockState
76         mutex   sync.Mutex
77 }
78
79 // NewBlockStateMap returns a newly allocated BlockStateMap.
80 func NewBlockStateMap() *BlockStateMap {
81         return &BlockStateMap{
82                 entries: make(map[arvados.SizedDigest]*BlockState),
83         }
84 }
85
86 // return a BlockState entry, allocating a new one if needed. (Private
87 // method: not goroutine-safe.)
88 func (bsm *BlockStateMap) get(blkid arvados.SizedDigest) *BlockState {
89         // TODO? Allocate BlockState structs a slice at a time,
90         // instead of one at a time.
91         blk := bsm.entries[blkid]
92         if blk == nil {
93                 blk = &BlockState{}
94                 bsm.entries[blkid] = blk
95         }
96         return blk
97 }
98
99 // Apply runs f on each entry in the map.
100 func (bsm *BlockStateMap) Apply(f func(arvados.SizedDigest, *BlockState)) {
101         bsm.mutex.Lock()
102         defer bsm.mutex.Unlock()
103
104         for blkid, blk := range bsm.entries {
105                 f(blkid, blk)
106         }
107 }
108
109 // AddReplicas updates the map to indicate that mnt has a replica of
110 // each block in idx.
111 func (bsm *BlockStateMap) AddReplicas(mnt *KeepMount, idx []arvados.KeepServiceIndexEntry) {
112         bsm.mutex.Lock()
113         defer bsm.mutex.Unlock()
114
115         for _, ent := range idx {
116                 bsm.get(ent.SizedDigest).addReplica(Replica{
117                         KeepMount: mnt,
118                         Mtime:     ent.Mtime,
119                 })
120         }
121 }
122
123 // IncreaseDesired updates the map to indicate the desired replication
124 // for the given blocks in the given storage class is at least n.
125 //
126 // If pdh is non-empty, it will be tracked and reported in the "lost
127 // blocks" report.
128 func (bsm *BlockStateMap) IncreaseDesired(pdh string, classes []string, n int, blocks []arvados.SizedDigest) {
129         bsm.mutex.Lock()
130         defer bsm.mutex.Unlock()
131
132         for _, blkid := range blocks {
133                 bsm.get(blkid).increaseDesired(pdh, classes, n)
134         }
135 }
136
137 // GetConfirmedReplication returns the replication level of the given
138 // blocks, considering only the specified storage classes.
139 //
140 // If len(classes)==0, returns the replication level without regard to
141 // storage classes.
142 //
143 // Safe to call concurrently with other calls to GetCurrent, but not
144 // with different BlockStateMap methods.
145 func (bsm *BlockStateMap) GetConfirmedReplication(blkids []arvados.SizedDigest, classes []string) int {
146         defaultClasses := map[string]bool{"default": true}
147         min := 0
148         for _, blkid := range blkids {
149                 total := 0
150                 perclass := make(map[string]int, len(classes))
151                 for _, c := range classes {
152                         perclass[c] = 0
153                 }
154                 bs, ok := bsm.entries[blkid]
155                 if !ok {
156                         return 0
157                 }
158                 for _, r := range bs.Replicas {
159                         total += r.KeepMount.Replication
160                         mntclasses := r.KeepMount.StorageClasses
161                         if len(mntclasses) == 0 {
162                                 mntclasses = defaultClasses
163                         }
164                         for c := range mntclasses {
165                                 n, ok := perclass[c]
166                                 if !ok {
167                                         // Don't care about this storage class
168                                         continue
169                                 }
170                                 perclass[c] = n + r.KeepMount.Replication
171                         }
172                 }
173                 if total == 0 {
174                         return 0
175                 }
176                 for _, n := range perclass {
177                         if n == 0 {
178                                 return 0
179                         }
180                         if n < min || min == 0 {
181                                 min = n
182                         }
183                 }
184                 if len(perclass) == 0 && (min == 0 || min > total) {
185                         min = total
186                 }
187         }
188         return min
189 }