diff --git a/quotacontrol.go b/quotacontrol.go index 15b5dc1..37da5fb 100644 --- a/quotacontrol.go +++ b/quotacontrol.go @@ -165,14 +165,11 @@ func (q qcHandler) GetAccessQuota(ctx context.Context, accessKey string, now tim Cycle: cycle, AccessKey: access, } - go func() { - // NOTE: we pass a fresh context here, as otherwise requests/other contexts which cancel - // above will cancel this goroutine and the AccessQuota will never be saved into cache. - err := q.quotaCache.SetAccessQuota(context.Background(), &record) - if err != nil { - q.log.With("error", err).Error("quotacontrol: failed to set access quota in cache") - } - }() + + if err := q.quotaCache.SetAccessQuota(ctx, &record); err != nil { + q.log.With("error", err).Error("quotacontrol: failed to set access quota in cache") + } + return &record, nil } @@ -401,23 +398,19 @@ func (q qcHandler) DisableAccessKey(ctx context.Context, accessKey string) (bool } func (q qcHandler) GetUserPermission(ctx context.Context, projectID uint64, userID string) (*proto.UserPermission, map[string]interface{}, error) { - userPerm, resourceAccess, err := q.permStore.GetUserPermission(ctx, projectID, userID) + perm, access, err := q.permStore.GetUserPermission(ctx, projectID, userID) if err != nil { - return userPerm, resourceAccess, proto.ErrUnauthorizedUser + return perm, access, proto.ErrUnauthorizedUser } - if userPerm != nil && *userPerm == proto.UserPermission_UNAUTHORIZED { - return userPerm, resourceAccess, proto.ErrUnauthorizedUser + if perm != nil && *perm == proto.UserPermission_UNAUTHORIZED { + return perm, access, proto.ErrUnauthorizedUser + } + + if err := q.permCache.SetUserPermission(ctx, projectID, userID, perm, access); err != nil { + q.log.With("error", err).Error("quotacontrol: failed to set user perm in cache") } - go func() { - // NOTE: we pass a fresh context here, as otherwise requests/other contexts which cancel - // above will cancel this goroutine and the AccessQuota will never be saved into cache. - err := q.permCache.SetUserPermission(context.Background(), projectID, userID, userPerm, resourceAccess) - if err != nil { - q.log.With("error", err).Error("quotacontrol: failed to set user perm in cache") - } - }() - return userPerm, resourceAccess, nil + return perm, access, nil } func (q qcHandler) updateAccessKey(ctx context.Context, access *proto.AccessKey) (*proto.AccessKey, error) { diff --git a/quotacontrol_test.go b/quotacontrol_test.go index 13d6ef3..a96d5a5 100644 --- a/quotacontrol_test.go +++ b/quotacontrol_test.go @@ -101,7 +101,6 @@ func TestMiddlewareUseAccessKey(t *testing.T) { } go func() { require.ErrorIs(t, server.ListenAndServe(), http.ErrServerClosed) }() defer server.Close() - time.Sleep(1 * time.Second) router := chi.NewRouter() @@ -124,7 +123,6 @@ func TestMiddlewareUseAccessKey(t *testing.T) { t.Run("WithAccessKey", func(t *testing.T) { go client.Run(context.Background()) - time.Sleep(1 * time.Second) ctx := middleware.WithTime(context.Background(), _Now) qc.notifications = make(map[uint64][]proto.EventType) @@ -204,7 +202,6 @@ func TestMiddlewareUseAccessKey(t *testing.T) { assert.NoError(t, err) go client.Run(context.Background()) - time.Sleep(1 * time.Second) ctx := middleware.WithTime(context.Background(), _Now) qc.notifications = make(map[uint64][]proto.EventType) @@ -223,7 +220,6 @@ func TestMiddlewareUseAccessKey(t *testing.T) { t.Run("PublicRateLimit", func(t *testing.T) { go client.Run(context.Background()) - time.Sleep(1 * time.Second) ctx := middleware.WithTime(context.Background(), _Now) @@ -336,7 +332,6 @@ func TestDefaultKey(t *testing.T) { } go func() { require.ErrorIs(t, server.ListenAndServe(), http.ErrServerClosed) }() defer server.Close() - time.Sleep(1 * time.Second) client := NewClient(wlog.With("client", "client"), *service, cfg)