bbgo_origin/pkg/bbgo/quota.go

package bbgo

import (
	"sync"

	"github.com/c9s/bbgo/pkg/fixedpoint"
)

type Quota struct {
	mu        sync.Mutex
	Available fixedpoint.Value
	Locked    fixedpoint.Value
}

// Add adds the fund to the available quota
func (q *Quota) Add(fund fixedpoint.Value) {
	q.mu.Lock()
	q.Available = q.Available.Add(fund)
	q.mu.Unlock()
}

// Lock locks the fund from the available quota
// returns true if the fund is locked successfully
// returns false if the fund is not enough
func (q *Quota) Lock(fund fixedpoint.Value) bool {
	if fund.Compare(q.Available) > 0 {
		return false
	}

	q.mu.Lock()
	q.Available = q.Available.Sub(fund)
	q.Locked = q.Locked.Add(fund)
	q.mu.Unlock()

	return true
}

// Commit commits the locked fund
func (q *Quota) Commit() {
	q.mu.Lock()
	q.Locked = fixedpoint.Zero
	q.mu.Unlock()
}

// Rollback rolls back the locked fund
// this will move the locked fund to the available quota
func (q *Quota) Rollback() {
	q.mu.Lock()
	q.Available = q.Available.Add(q.Locked)
	q.Locked = fixedpoint.Zero
	q.mu.Unlock()
}

func (q *Quota) String() string {
	q.mu.Lock()
	defer q.mu.Unlock()

	return q.Locked.String() + "/" + q.Available.String()
}

// QuotaTransaction is a transactional quota manager
type QuotaTransaction struct {
	mu         sync.Mutex
	BaseAsset  Quota
	QuoteAsset Quota
}

// Commit commits the transaction
func (m *QuotaTransaction) Commit() bool {
	m.mu.Lock()
	m.BaseAsset.Commit()
	m.QuoteAsset.Commit()
	m.mu.Unlock()
	return true
}

// Rollback rolls back the transaction
func (m *QuotaTransaction) Rollback() bool {
	m.mu.Lock()
	m.BaseAsset.Rollback()
	m.QuoteAsset.Rollback()
	m.mu.Unlock()
	return true
}