storj/satellite/audit/reservoir.go
dlamarmorgan b3cea3d1b6 satellite/audit: account for piece size during audit reservoir sampling
Treat the piece size as a weight, and perform weighted reservoir sampling as given in Algorithm A-Chao (https://en.wikipedia.org/wiki/Reservoir_sampling#Algorithm_A-Chao)

Change-Id: I299d0026d9e02d03b3d2130b0f32192928e6e326
2021-12-01 18:17:52 +00:00

82 lines
2.3 KiB
Go

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.
package audit
import (
"math/rand"
"time"
"storj.io/common/uuid"
"storj.io/storj/satellite/metabase"
"storj.io/storj/satellite/metabase/segmentloop"
)
const maxReservoirSize = 3
// Reservoir holds a certain number of segments to reflect a random sample.
type Reservoir struct {
Segments [maxReservoirSize]Segment
size int8
index int64
wSum int64
}
// NewReservoir instantiates a Reservoir.
func NewReservoir(size int) *Reservoir {
if size < 1 {
size = 1
} else if size > maxReservoirSize {
size = maxReservoirSize
}
return &Reservoir{
size: int8(size),
index: 0,
}
}
// Sample makes sure that for every segment in metainfo from index i=size..n-1,
// compute the relative weight based on segment size, and pick a random floating
// point number r = rand(0..1), and if r < the relative weight of the segment,
// select uniformly a random segment reservoir.Segments[rand(0..i)] to replace with
// segment. See https://en.wikipedia.org/wiki/Reservoir_sampling#Algorithm_A-Chao
// for the algorithm used.
func (reservoir *Reservoir) Sample(r *rand.Rand, segment Segment) {
if reservoir.index < int64(reservoir.size) {
reservoir.Segments[reservoir.index] = segment
reservoir.wSum += int64(segment.EncryptedSize)
} else {
reservoir.wSum += int64(segment.EncryptedSize)
p := float64(segment.EncryptedSize) / float64(reservoir.wSum)
random := r.Float64()
if random < p {
index := r.Int31n(int32(reservoir.size))
reservoir.Segments[index] = segment
}
}
reservoir.index++
}
// Segment is a segment to audit.
type Segment struct {
StreamID uuid.UUID
Position metabase.SegmentPosition
ExpiresAt *time.Time
EncryptedSize int32 // size of the whole segment (not a piece)
}
// NewSegment creates a new segment to audit from a metainfo loop segment.
func NewSegment(loopSegment *segmentloop.Segment) Segment {
return Segment{
StreamID: loopSegment.StreamID,
Position: loopSegment.Position,
ExpiresAt: loopSegment.ExpiresAt,
EncryptedSize: loopSegment.EncryptedSize,
}
}
// Expired checks if segment is expired relative to now.
func (segment *Segment) Expired(now time.Time) bool {
return segment.ExpiresAt != nil && segment.ExpiresAt.Before(now)
}