Skip to content

Commit 84e0e9a

Browse files
cheggaaacheggaaa
authored
Merge pull request #58 from anyproto/GO-2837-space-limits
GO-2837 space limits
2 parents 60be868 + 46818f2 commit 84e0e9a

File tree

16 files changed

+441
-1109
lines changed

16 files changed

+441
-1109
lines changed

‎accountlimit/accountlimit.go‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
package accountlimit
2+
3+
import (
4+
"context"
5+
"errors"
6+
"time"
7+
8+
"github.com/anyproto/any-sync/app"
9+
"github.com/anyproto/any-sync/app/logger"
10+
"github.com/anyproto/any-sync/commonfile/fileproto"
11+
"github.com/anyproto/any-sync/net/pool"
12+
"github.com/anyproto/any-sync/nodeconf"
13+
"go.mongodb.org/mongo-driver/bson"
14+
"go.mongodb.org/mongo-driver/mongo"
15+
"go.mongodb.org/mongo-driver/mongo/options"
16+
"storj.io/drpc"
17+
18+
"github.com/anyproto/any-sync-coordinator/db"
19+
"github.com/anyproto/any-sync-coordinator/spacestatus"
20+
)
21+
22+
func New() AccountLimit {
23+
return &accountLimit{}
24+
}
25+
26+
const CName = "coordinator.accountLimit"
27+
28+
const collName = "accountLimit"
29+
30+
var log = logger.NewNamed(CName)
31+
32+
type configGetter interface {
33+
GetAccountLimit() SpaceLimits
34+
}
35+
36+
type SpaceLimits struct {
37+
SpaceMembersRead uint32 `yaml:"spaceMembersRead" bson:"spaceMembersRead"`
38+
SpaceMembersWrite uint32 `yaml:"spaceMembersWrite" bson:"spaceMembersWrite"`
39+
}
40+
41+
type Limits struct {
42+
Identity string `bson:"_id"`
43+
Reason string `bson:"reason"`
44+
FileStorageBytes uint64 `bson:"fileStorageBytes"`
45+
SpaceMembersRead uint32 `bson:"spaceMembersRead"`
46+
SpaceMembersWrite uint32 `bson:"spaceMembersWrite"`
47+
UpdatedTime time.Time `bson:"updatedTime"`
48+
}
49+
50+
type AccountLimit interface {
51+
SetLimits(ctx context.Context, limits Limits) (err error)
52+
GetLimits(ctx context.Context, identity string) (limits Limits, err error)
53+
GetLimitsBySpace(ctx context.Context, spaceId string) (limits SpaceLimits, err error)
54+
app.Component
55+
}
56+
type accountLimit struct {
57+
pool pool.Pool
58+
nodeConf nodeconf.Service
59+
coll *mongo.Collection
60+
spaceStatus spacestatus.SpaceStatus
61+
defaultLimits SpaceLimits
62+
}
63+
64+
func (al *accountLimit) Init(a *app.App) (err error) {
65+
al.pool = a.MustComponent(pool.CName).(pool.Pool)
66+
al.nodeConf = a.MustComponent(nodeconf.CName).(nodeconf.Service)
67+
al.coll = a.MustComponent(db.CName).(db.Database).Db().Collection(collName)
68+
al.spaceStatus = a.MustComponent(spacestatus.CName).(spacestatus.SpaceStatus)
69+
al.defaultLimits = a.MustComponent("config").(configGetter).GetAccountLimit()
70+
return nil
71+
}
72+
73+
func (al *accountLimit) Name() (name string) {
74+
return CName
75+
}
76+
77+
func (al *accountLimit) SetLimits(ctx context.Context, limits Limits) (err error) {
78+
limits.UpdatedTime = time.Now()
79+
80+
if err = al.updateFileLimits(ctx, limits); err != nil {
81+
return
82+
}
83+
84+
return al.updateDbLimits(ctx, limits)
85+
}
86+
87+
func (al *accountLimit) updateDbLimits(ctx context.Context, limits Limits) (err error) {
88+
_, err = al.coll.UpdateOne(
89+
ctx,
90+
bson.D{{"_id", limits.Identity}},
91+
bson.D{{"$set", limits}},
92+
options.Update().SetUpsert(true),
93+
)
94+
return
95+
}
96+
97+
func (al *accountLimit) updateFileLimits(ctx context.Context, limits Limits) (err error) {
98+
filePeer, err := al.pool.GetOneOf(ctx, al.nodeConf.FilePeers())
99+
if err != nil {
100+
return
101+
}
102+
return filePeer.DoDrpc(ctx, func(conn drpc.Conn) error {
103+
_, err := fileproto.NewDRPCFileClient(conn).AccountLimitSet(ctx, &fileproto.AccountLimitSetRequest{
104+
Identity: limits.Identity,
105+
Limit: limits.FileStorageBytes,
106+
})
107+
return err
108+
})
109+
}
110+
111+
func (al *accountLimit) GetLimits(ctx context.Context, identity string) (limits Limits, err error) {
112+
err = al.coll.FindOne(ctx, bson.D{{"_id", identity}}).Decode(&limits)
113+
if err == nil || !errors.Is(err, mongo.ErrNoDocuments) {
114+
return
115+
}
116+
// default limit
117+
return Limits{
118+
Identity: identity,
119+
SpaceMembersRead: al.defaultLimits.SpaceMembersRead,
120+
SpaceMembersWrite: al.defaultLimits.SpaceMembersWrite,
121+
UpdatedTime: time.Now(),
122+
}, nil
123+
}
124+
125+
func (al *accountLimit) GetLimitsBySpace(ctx context.Context, spaceId string) (sLimits SpaceLimits, err error) {
126+
entry, err := al.spaceStatus.Status(ctx, spaceId)
127+
if err != nil {
128+
return
129+
}
130+
131+
// return 1-1 to personal and tech spaces
132+
switch entry.Type {
133+
case spacestatus.SpaceTypePersonal, spacestatus.SpaceTypeTech:
134+
return SpaceLimits{
135+
SpaceMembersRead: 1,
136+
SpaceMembersWrite: 1,
137+
}, nil
138+
}
139+
140+
limits, err := al.GetLimits(ctx, entry.Identity)
141+
if err != nil {
142+
return
143+
}
144+
return SpaceLimits{
145+
SpaceMembersRead: limits.SpaceMembersRead,
146+
SpaceMembersWrite: limits.SpaceMembersWrite,
147+
}, nil
148+
}

‎accountlimit/accountlimit_test.go‎

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
package accountlimit
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/anyproto/any-sync/app"
9+
"github.com/anyproto/any-sync/net/peer/mock_peer"
10+
"github.com/anyproto/any-sync/net/rpc/rpctest"
11+
"github.com/anyproto/any-sync/nodeconf"
12+
"github.com/anyproto/any-sync/nodeconf/mock_nodeconf"
13+
"github.com/stretchr/testify/assert"
14+
"github.com/stretchr/testify/require"
15+
"go.uber.org/mock/gomock"
16+
17+
"github.com/anyproto/any-sync-coordinator/db"
18+
"github.com/anyproto/any-sync-coordinator/spacestatus"
19+
"github.com/anyproto/any-sync-coordinator/spacestatus/mock_spacestatus"
20+
)
21+
22+
func TestAccountLimit_SetLimits(t *testing.T) {
23+
fx := newFixture(t)
24+
defer fx.finish(t)
25+
26+
fx.nodeConf.EXPECT().FilePeers().Return([]string{"filePeer"}).Times(2)
27+
peer := mock_peer.NewMockPeer(fx.ctrl)
28+
peer.EXPECT().Id().Return("filePeer").AnyTimes()
29+
peer.EXPECT().DoDrpc(ctx, gomock.Any()).Times(2)
30+
fx.pool.AddPeer(ctx, peer)
31+
32+
limits := Limits{
33+
Identity: "123",
34+
SpaceMembersRead: 100,
35+
SpaceMembersWrite: 200,
36+
}
37+
38+
// set
39+
require.NoError(t, fx.SetLimits(ctx, limits))
40+
result, err := fx.GetLimits(ctx, "123")
41+
require.NoError(t, err)
42+
result.UpdatedTime = time.Time{}
43+
assert.Equal(t, limits, result)
44+
45+
// update
46+
limits.SpaceMembersRead = 1000
47+
limits.SpaceMembersWrite = 2000
48+
limits.Reason = "upsert"
49+
require.NoError(t, fx.SetLimits(ctx, limits))
50+
result, err = fx.GetLimits(ctx, "123")
51+
require.NoError(t, err)
52+
result.UpdatedTime = time.Time{}
53+
assert.Equal(t, limits, result)
54+
55+
}
56+
57+
func TestAccountLimit_GetLimits(t *testing.T) {
58+
fx := newFixture(t)
59+
defer fx.finish(t)
60+
61+
// get default limits
62+
limits, err := fx.GetLimits(ctx, "default")
63+
require.NoError(t, err)
64+
assert.Equal(t, uint32(5), limits.SpaceMembersWrite)
65+
assert.Equal(t, uint32(10), limits.SpaceMembersRead)
66+
}
67+
68+
func TestAccountLimit_GetLimitsBySpace(t *testing.T) {
69+
var (
70+
spaceId = "spaceId"
71+
identity = "identity"
72+
)
73+
t.Run("personal", func(t *testing.T) {
74+
fx := newFixture(t)
75+
defer fx.finish(t)
76+
77+
fx.spaceStatus.EXPECT().Status(ctx, spaceId).Return(spacestatus.StatusEntry{
78+
SpaceId: spaceId,
79+
Identity: identity,
80+
Type: spacestatus.SpaceTypePersonal,
81+
}, nil)
82+
83+
limits, err := fx.GetLimitsBySpace(ctx, spaceId)
84+
require.NoError(t, err)
85+
assert.Equal(t, SpaceLimits{1, 1}, limits)
86+
})
87+
88+
t.Run("regular", func(t *testing.T) {
89+
fx := newFixture(t)
90+
defer fx.finish(t)
91+
92+
fx.spaceStatus.EXPECT().Status(ctx, spaceId).Return(spacestatus.StatusEntry{
93+
SpaceId: spaceId,
94+
Identity: identity,
95+
Type: spacestatus.SpaceTypeRegular,
96+
}, nil)
97+
98+
limits, err := fx.GetLimitsBySpace(ctx, spaceId)
99+
require.NoError(t, err)
100+
assert.Equal(t, SpaceLimits{10, 5}, limits)
101+
})
102+
103+
}
104+
105+
var ctx = context.Background()
106+
107+
func newFixture(t *testing.T) *fixture {
108+
ctrl := gomock.NewController(t)
109+
fx := &fixture{
110+
AccountLimit: New(),
111+
nodeConf: mock_nodeconf.NewMockService(ctrl),
112+
pool: rpctest.NewTestPool(),
113+
spaceStatus: mock_spacestatus.NewMockSpaceStatus(ctrl),
114+
a: new(app.App),
115+
ctrl: ctrl,
116+
}
117+
118+
fx.nodeConf.EXPECT().Init(gomock.Any()).AnyTimes()
119+
fx.nodeConf.EXPECT().Name().Return(nodeconf.CName).AnyTimes()
120+
fx.nodeConf.EXPECT().Run(gomock.Any()).AnyTimes()
121+
fx.nodeConf.EXPECT().Close(gomock.Any()).AnyTimes()
122+
123+
fx.spaceStatus.EXPECT().Init(gomock.Any()).AnyTimes()
124+
fx.spaceStatus.EXPECT().Name().Return(spacestatus.CName).AnyTimes()
125+
fx.spaceStatus.EXPECT().Run(gomock.Any()).AnyTimes()
126+
fx.spaceStatus.EXPECT().Close(gomock.Any()).AnyTimes()
127+
128+
fx.a.Register(db.New()).
129+
Register(fx.AccountLimit).
130+
Register(fx.nodeConf).
131+
Register(fx.spaceStatus).
132+
Register(fx.pool).
133+
Register(&testConfig{})
134+
135+
require.NoError(t, fx.a.Start(ctx))
136+
_ = fx.a.MustComponent(db.CName).(db.Database).Db().Collection(collName).Drop(ctx)
137+
return fx
138+
}
139+
140+
type fixture struct {
141+
AccountLimit
142+
a *app.App
143+
ctrl *gomock.Controller
144+
nodeConf *mock_nodeconf.MockService
145+
pool *rpctest.TestPool
146+
spaceStatus *mock_spacestatus.MockSpaceStatus
147+
}
148+
149+
func (fx *fixture) finish(t *testing.T) {
150+
require.NoError(t, fx.a.Close(ctx))
151+
}
152+
153+
type testConfig struct {
154+
}
155+
156+
func (c *testConfig) Init(_ *app.App) error { return nil }
157+
func (c *testConfig) Name() string { return "config" }
158+
159+
func (c *testConfig) GetMongo() db.Mongo {
160+
return db.Mongo{
161+
Connect: "mongodb://localhost:27017",
162+
Database: "coordinator_unittest",
163+
}
164+
}
165+
166+
func (c *testConfig) GetAccountLimit() SpaceLimits {
167+
return SpaceLimits{
168+
SpaceMembersWrite: 5,
169+
SpaceMembersRead: 10,
170+
}
171+
}

0 commit comments

Comments
 (0)