mirror of
https://github.com/c9s/bbgo.git
synced 2024-09-20 08:11:08 +00:00
85 lines
1.6 KiB
Go
85 lines
1.6 KiB
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
|
|
}
|