Merge branch '13340-wb-keep-links'
[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 main
6
7 import (
8         "sync"
9
10         "git.curoverse.com/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         Replicas []Replica
27         Desired  map[string]int
28         // TODO: Support combinations of classes ("private + durable")
29         // by replacing the map[string]int with a map[*[]string]int
30         // here, where the map keys come from a pool of semantically
31         // distinct class combinations.
32         //
33         // TODO: Use a pool of semantically distinct Desired maps to
34         // conserve memory (typically there are far more BlockState
35         // objects in memory than distinct Desired profiles).
36 }
37
38 var defaultClasses = []string{"default"}
39
40 func (bs *BlockState) addReplica(r Replica) {
41         bs.Replicas = append(bs.Replicas, r)
42 }
43
44 func (bs *BlockState) increaseDesired(classes []string, n int) {
45         if len(classes) == 0 {
46                 classes = defaultClasses
47         }
48         for _, class := range classes {
49                 if bs.Desired == nil {
50                         bs.Desired = map[string]int{class: n}
51                 } else if d, ok := bs.Desired[class]; !ok || d < n {
52                         bs.Desired[class] = n
53                 }
54         }
55 }
56
57 // BlockStateMap is a goroutine-safe wrapper around a
58 // map[arvados.SizedDigest]*BlockState.
59 type BlockStateMap struct {
60         entries map[arvados.SizedDigest]*BlockState
61         mutex   sync.Mutex
62 }
63
64 // NewBlockStateMap returns a newly allocated BlockStateMap.
65 func NewBlockStateMap() *BlockStateMap {
66         return &BlockStateMap{
67                 entries: make(map[arvados.SizedDigest]*BlockState),
68         }
69 }
70
71 // return a BlockState entry, allocating a new one if needed. (Private
72 // method: not goroutine-safe.)
73 func (bsm *BlockStateMap) get(blkid arvados.SizedDigest) *BlockState {
74         // TODO? Allocate BlockState structs a slice at a time,
75         // instead of one at a time.
76         blk := bsm.entries[blkid]
77         if blk == nil {
78                 blk = &BlockState{}
79                 bsm.entries[blkid] = blk
80         }
81         return blk
82 }
83
84 // Apply runs f on each entry in the map.
85 func (bsm *BlockStateMap) Apply(f func(arvados.SizedDigest, *BlockState)) {
86         bsm.mutex.Lock()
87         defer bsm.mutex.Unlock()
88
89         for blkid, blk := range bsm.entries {
90                 f(blkid, blk)
91         }
92 }
93
94 // AddReplicas updates the map to indicate that mnt has a replica of
95 // each block in idx.
96 func (bsm *BlockStateMap) AddReplicas(mnt *KeepMount, idx []arvados.KeepServiceIndexEntry) {
97         bsm.mutex.Lock()
98         defer bsm.mutex.Unlock()
99
100         for _, ent := range idx {
101                 bsm.get(ent.SizedDigest).addReplica(Replica{
102                         KeepMount: mnt,
103                         Mtime:     ent.Mtime,
104                 })
105         }
106 }
107
108 // IncreaseDesired updates the map to indicate the desired replication
109 // for the given blocks in the given storage class is at least n.
110 func (bsm *BlockStateMap) IncreaseDesired(classes []string, n int, blocks []arvados.SizedDigest) {
111         bsm.mutex.Lock()
112         defer bsm.mutex.Unlock()
113
114         for _, blkid := range blocks {
115                 bsm.get(blkid).increaseDesired(classes, n)
116         }
117 }