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         *KeepService
  18         Mtime int64
  19 }
  20
  21 // BlockState indicates the number of desired replicas (according to
  22 // the collections we know about) and the replicas actually stored
  23 // (according to the keepstore indexes we know about).
  24 type BlockState struct {
  25         Replicas []Replica
  26         Desired  int
  27 }
  28
  29 func (bs *BlockState) addReplica(r Replica) {
  30         bs.Replicas = append(bs.Replicas, r)
  31 }
  32
  33 func (bs *BlockState) increaseDesired(n int) {
  34         if bs.Desired < n {
  35                 bs.Desired = n
  36         }
  37 }
  38
  39 // BlockStateMap is a goroutine-safe wrapper around a
  40 // map[arvados.SizedDigest]*BlockState.
  41 type BlockStateMap struct {
  42         entries map[arvados.SizedDigest]*BlockState
  43         mutex   sync.Mutex
  44 }
  45
  46 // NewBlockStateMap returns a newly allocated BlockStateMap.
  47 func NewBlockStateMap() *BlockStateMap {
  48         return &BlockStateMap{
  49                 entries: make(map[arvados.SizedDigest]*BlockState),
  50         }
  51 }
  52
  53 // return a BlockState entry, allocating a new one if needed. (Private
  54 // method: not goroutine-safe.)
  55 func (bsm *BlockStateMap) get(blkid arvados.SizedDigest) *BlockState {
  56         // TODO? Allocate BlockState structs a slice at a time,
  57         // instead of one at a time.
  58         blk := bsm.entries[blkid]
  59         if blk == nil {
  60                 blk = &BlockState{}
  61                 bsm.entries[blkid] = blk
  62         }
  63         return blk
  64 }
  65
  66 // Apply runs f on each entry in the map.
  67 func (bsm *BlockStateMap) Apply(f func(arvados.SizedDigest, *BlockState)) {
  68         bsm.mutex.Lock()
  69         defer bsm.mutex.Unlock()
  70
  71         for blkid, blk := range bsm.entries {
  72                 f(blkid, blk)
  73         }
  74 }
  75
  76 // AddReplicas updates the map to indicate srv has a replica of each
  77 // block in idx.
  78 func (bsm *BlockStateMap) AddReplicas(srv *KeepService, idx []arvados.KeepServiceIndexEntry) {
  79         bsm.mutex.Lock()
  80         defer bsm.mutex.Unlock()
  81
  82         for _, ent := range idx {
  83                 bsm.get(ent.SizedDigest).addReplica(Replica{
  84                         KeepService: srv,
  85                         Mtime:       ent.Mtime,
  86                 })
  87         }
  88 }
  89
  90 // IncreaseDesired updates the map to indicate the desired replication
  91 // for the given blocks is at least n.
  92 func (bsm *BlockStateMap) IncreaseDesired(n int, blocks []arvados.SizedDigest) {
  93         bsm.mutex.Lock()
  94         defer bsm.mutex.Unlock()
  95
  96         for _, blkid := range blocks {
  97                 bsm.get(blkid).increaseDesired(n)
  98         }
  99 }