github.com/ethersphere/bee/v2@v2.2.0/pkg/storer/mock/forgetting.go (about)

     1  // Copyright 2023 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package mockstorer
     6  
     7  import (
     8  	"context"
     9  	"sync"
    10  	"sync/atomic"
    11  	"time"
    12  
    13  	storage "github.com/ethersphere/bee/v2/pkg/storage"
    14  	"github.com/ethersphere/bee/v2/pkg/swarm"
    15  )
    16  
    17  type DelayedStore struct {
    18  	storage.ChunkStore
    19  	cache map[string]time.Duration
    20  	mu    sync.Mutex
    21  }
    22  
    23  func NewDelayedStore(s storage.ChunkStore) *DelayedStore {
    24  	return &DelayedStore{
    25  		ChunkStore: s,
    26  		cache:      make(map[string]time.Duration),
    27  	}
    28  }
    29  
    30  func (d *DelayedStore) Delay(addr swarm.Address, delay time.Duration) {
    31  	d.mu.Lock()
    32  	defer d.mu.Unlock()
    33  	d.cache[addr.String()] = delay
    34  }
    35  
    36  func (d *DelayedStore) Get(ctx context.Context, addr swarm.Address) (ch swarm.Chunk, err error) {
    37  	d.mu.Lock()
    38  	delay, ok := d.cache[addr.String()]
    39  	if ok && delay > 0 {
    40  		delete(d.cache, addr.String())
    41  		d.mu.Unlock()
    42  		select {
    43  		case <-time.After(delay):
    44  		case <-ctx.Done():
    45  			return nil, ctx.Err()
    46  		}
    47  	} else {
    48  		d.mu.Unlock()
    49  	}
    50  	return d.ChunkStore.Get(ctx, addr)
    51  }
    52  
    53  type ForgettingStore struct {
    54  	storage.ChunkStore
    55  	record atomic.Bool
    56  	mu     sync.Mutex
    57  	n      atomic.Int64
    58  	missed map[string]struct{}
    59  }
    60  
    61  func NewForgettingStore(s storage.ChunkStore) *ForgettingStore {
    62  	return &ForgettingStore{ChunkStore: s, missed: make(map[string]struct{})}
    63  }
    64  
    65  func (f *ForgettingStore) Stored() int64 {
    66  	return f.n.Load()
    67  }
    68  
    69  func (f *ForgettingStore) Record() {
    70  	f.mu.Lock()
    71  	defer f.mu.Unlock()
    72  	f.record.Store(true)
    73  }
    74  
    75  func (f *ForgettingStore) Unrecord() {
    76  	f.mu.Lock()
    77  	defer f.mu.Unlock()
    78  	f.record.Store(false)
    79  }
    80  
    81  func (f *ForgettingStore) Miss(addr swarm.Address) {
    82  	f.mu.Lock()
    83  	defer f.mu.Unlock()
    84  	f.miss(addr)
    85  }
    86  
    87  func (f *ForgettingStore) Unmiss(addr swarm.Address) {
    88  	f.mu.Lock()
    89  	defer f.mu.Unlock()
    90  	f.unmiss(addr)
    91  }
    92  
    93  func (f *ForgettingStore) miss(addr swarm.Address) {
    94  	f.missed[addr.String()] = struct{}{}
    95  }
    96  
    97  func (f *ForgettingStore) unmiss(addr swarm.Address) {
    98  	delete(f.missed, addr.String())
    99  }
   100  
   101  func (f *ForgettingStore) isMiss(addr swarm.Address) bool {
   102  	_, ok := f.missed[addr.String()]
   103  	return ok
   104  }
   105  
   106  func (f *ForgettingStore) Reset() {
   107  	f.mu.Lock()
   108  	defer f.mu.Unlock()
   109  	f.missed = make(map[string]struct{})
   110  }
   111  
   112  func (f *ForgettingStore) Missed() int {
   113  	f.mu.Lock()
   114  	defer f.mu.Unlock()
   115  	return len(f.missed)
   116  }
   117  
   118  // Get implements the ChunkStore interface.
   119  // if in recording phase, record the chunk address as miss and returns Get on the embedded store
   120  // if in forgetting phase, returns ErrNotFound if the chunk address is recorded as miss
   121  func (f *ForgettingStore) Get(ctx context.Context, addr swarm.Address) (ch swarm.Chunk, err error) {
   122  	f.mu.Lock()
   123  	defer f.mu.Unlock()
   124  	if f.record.Load() {
   125  		f.miss(addr)
   126  	} else if f.isMiss(addr) {
   127  		return nil, storage.ErrNotFound
   128  	}
   129  	return f.ChunkStore.Get(ctx, addr)
   130  }
   131  
   132  // Put implements the ChunkStore interface.
   133  func (f *ForgettingStore) Put(ctx context.Context, ch swarm.Chunk) (err error) {
   134  	f.n.Add(1)
   135  	if !f.record.Load() {
   136  		f.Unmiss(ch.Address())
   137  	}
   138  	return f.ChunkStore.Put(ctx, ch)
   139  }