bbgo_origin/pkg/bbgo/quota.go

85 lines
1.6 KiB
Go
Raw Normal View History

2020-11-23 08:47:36 +00:00
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
2020-11-23 08:47:36 +00:00
func (q *Quota) Add(fund fixedpoint.Value) {
q.mu.Lock()
q.Available = q.Available.Add(fund)
2020-11-23 08:47:36 +00:00
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
2020-11-23 08:47:36 +00:00
func (q *Quota) Lock(fund fixedpoint.Value) bool {
if fund.Compare(q.Available) > 0 {
2020-11-23 08:47:36 +00:00
return false
}
q.mu.Lock()
q.Available = q.Available.Sub(fund)
q.Locked = q.Locked.Add(fund)
2020-11-23 08:47:36 +00:00
q.mu.Unlock()
return true
}
// Commit commits the locked fund
2020-11-23 08:47:36 +00:00
func (q *Quota) Commit() {
q.mu.Lock()
q.Locked = fixedpoint.Zero
2020-11-23 08:47:36 +00:00
q.mu.Unlock()
}
// Rollback rolls back the locked fund
// this will move the locked fund to the available quota
2020-11-23 08:47:36 +00:00
func (q *Quota) Rollback() {
q.mu.Lock()
q.Available = q.Available.Add(q.Locked)
q.Locked = fixedpoint.Zero
2020-11-23 08:47:36 +00:00
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
2020-11-23 08:47:36 +00:00
type QuotaTransaction struct {
mu sync.Mutex
BaseAsset Quota
QuoteAsset Quota
}
// Commit commits the transaction
2020-11-23 08:47:36 +00:00
func (m *QuotaTransaction) Commit() bool {
m.mu.Lock()
m.BaseAsset.Commit()
m.QuoteAsset.Commit()
m.mu.Unlock()
return true
}
// Rollback rolls back the transaction
2020-11-23 08:47:36 +00:00
func (m *QuotaTransaction) Rollback() bool {
m.mu.Lock()
m.BaseAsset.Rollback()
m.QuoteAsset.Rollback()
m.mu.Unlock()
return true
}