operations: add method to real time account server side copy
Before this change server side copies would show at 0% until they were done then show at 100%. With support from the backend, server side copies can now be accounted in real time. This will only work for backends which have been modified and themselves get feedback about how copies are going. If the transfer fails, the bytes accounted will be reversed.
This commit is contained in:
@@ -0,0 +1,75 @@
|
||||
// Package transferaccounter provides utilities for accounting server side transfers.
|
||||
package transferaccounter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// Context key type for accounter
|
||||
type accounterContextKeyType struct{}
|
||||
|
||||
// Context key for accounter
|
||||
var accounterContextKey = accounterContextKeyType{}
|
||||
|
||||
// TransferAccounter is used to account server side and other transfers.
|
||||
type TransferAccounter struct {
|
||||
add func(n int64)
|
||||
total atomic.Int64
|
||||
started bool
|
||||
}
|
||||
|
||||
// New creates a TransferAccounter using the add function passed in.
|
||||
//
|
||||
// Note that the add function should be goroutine safe.
|
||||
//
|
||||
// It adds the new TransferAccounter to the context.
|
||||
func New(ctx context.Context, add func(n int64)) (context.Context, *TransferAccounter) {
|
||||
ta := &TransferAccounter{
|
||||
add: add,
|
||||
}
|
||||
newCtx := context.WithValue(ctx, accounterContextKey, ta)
|
||||
return newCtx, ta
|
||||
}
|
||||
|
||||
// Start the transfer. Call this before calling Add().
|
||||
func (ta *TransferAccounter) Start() {
|
||||
ta.started = true
|
||||
}
|
||||
|
||||
// Started returns if the transfer has had Start() called or not.
|
||||
func (ta *TransferAccounter) Started() bool {
|
||||
return ta.started
|
||||
}
|
||||
|
||||
// Add n bytes to the transfer
|
||||
func (ta *TransferAccounter) Add(n int64) {
|
||||
ta.add(n)
|
||||
ta.total.Add(n)
|
||||
}
|
||||
|
||||
// Reset reverses out all accounted stats if Started() has been called
|
||||
func (ta *TransferAccounter) Reset() {
|
||||
if ta.started {
|
||||
ta.Add(-ta.total.Load())
|
||||
}
|
||||
}
|
||||
|
||||
// A transfer accounter which does nothing
|
||||
var nullAccounter = &TransferAccounter{
|
||||
add: func(n int64) {},
|
||||
}
|
||||
|
||||
// Get returns a *TransferAccounter from the ctx.
|
||||
//
|
||||
// If none is found it will return a dummy one to keep the code simple.
|
||||
func Get(ctx context.Context) *TransferAccounter {
|
||||
if ctx == nil {
|
||||
return nullAccounter
|
||||
}
|
||||
c := ctx.Value(accounterContextKey)
|
||||
if c == nil {
|
||||
return nullAccounter
|
||||
}
|
||||
return c.(*TransferAccounter)
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package transferaccounter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNew(t *testing.T) {
|
||||
// Dummy add function
|
||||
var totalBytes int64
|
||||
addFn := func(n int64) {
|
||||
totalBytes += n
|
||||
}
|
||||
|
||||
// Create the accounter
|
||||
ctx := context.Background()
|
||||
_, ta := New(ctx, addFn)
|
||||
|
||||
// Verify object creation
|
||||
require.NotNil(t, ta)
|
||||
assert.False(t, ta.Started(), "New accounter should not be started by default")
|
||||
|
||||
// Test Start()
|
||||
ta.Start()
|
||||
assert.True(t, ta.Started(), "Accounter should be started after calling Start()")
|
||||
|
||||
// Test Add() logic
|
||||
ta.Add(100)
|
||||
ta.Add(50)
|
||||
assert.Equal(t, int64(150), totalBytes, "The add function should have been called with cumulative values")
|
||||
assert.Equal(t, int64(150), ta.total.Load(), "Internal counter did not count")
|
||||
|
||||
// Test Reset() logic
|
||||
ta.Reset()
|
||||
assert.Equal(t, int64(0), totalBytes, "The Reset function failed")
|
||||
assert.Equal(t, int64(0), ta.total.Load(), "Internal counter did not reset")
|
||||
|
||||
}
|
||||
|
||||
func TestGet(t *testing.T) {
|
||||
t.Run("Retrieve existing accounter", func(t *testing.T) {
|
||||
// Create a specific accounter to identify later
|
||||
expectedTotal := int64(0)
|
||||
ctx, originalTa := New(context.Background(), func(n int64) { expectedTotal += n })
|
||||
|
||||
// Retrieve it
|
||||
retrievedTa := Get(ctx)
|
||||
|
||||
// Assert it is the exact same pointer
|
||||
assert.Equal(t, originalTa, retrievedTa)
|
||||
|
||||
// Verify functionality passes through
|
||||
retrievedTa.Add(10)
|
||||
assert.Equal(t, int64(10), expectedTotal)
|
||||
})
|
||||
|
||||
t.Run("Context does not contain accounter", func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ta := Get(ctx)
|
||||
|
||||
assert.NotNil(t, ta, "Get should never return nil")
|
||||
assert.Equal(t, nullAccounter, ta, "Should return the global nullAccounter")
|
||||
})
|
||||
|
||||
t.Run("Context is nil", func(t *testing.T) {
|
||||
ta := Get(nil) //nolint:staticcheck // we want to test this
|
||||
|
||||
assert.NotNil(t, ta, "Get should never return nil")
|
||||
assert.Equal(t, nullAccounter, ta, "Should return the global nullAccounter")
|
||||
})
|
||||
}
|
||||
|
||||
func TestNullAccounterBehavior(t *testing.T) {
|
||||
// Ensure the null accounter (returned when context is missing/nil)
|
||||
// can be called without panicking.
|
||||
ta := Get(nil) //nolint:staticcheck // we want to test this
|
||||
|
||||
assert.NotPanics(t, func() {
|
||||
ta.Start()
|
||||
})
|
||||
|
||||
// Even after start, it acts as a valid object
|
||||
assert.True(t, ta.Started())
|
||||
|
||||
assert.NotPanics(t, func() {
|
||||
ta.Add(1000)
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user