1// Copyright 2016 The etcd Authors 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package v3rpc 16 17import ( 18 "github.com/coreos/etcd/etcdserver" 19 "github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes" 20 pb "github.com/coreos/etcd/etcdserver/etcdserverpb" 21 "github.com/coreos/etcd/pkg/types" 22 "golang.org/x/net/context" 23) 24 25type quotaKVServer struct { 26 pb.KVServer 27 qa quotaAlarmer 28} 29 30type quotaAlarmer struct { 31 q etcdserver.Quota 32 a Alarmer 33 id types.ID 34} 35 36// check whether request satisfies the quota. If there is not enough space, 37// ignore request and raise the free space alarm. 38func (qa *quotaAlarmer) check(ctx context.Context, r interface{}) error { 39 if qa.q.Available(r) { 40 return nil 41 } 42 req := &pb.AlarmRequest{ 43 MemberID: uint64(qa.id), 44 Action: pb.AlarmRequest_ACTIVATE, 45 Alarm: pb.AlarmType_NOSPACE, 46 } 47 qa.a.Alarm(ctx, req) 48 return rpctypes.ErrGRPCNoSpace 49} 50 51func NewQuotaKVServer(s *etcdserver.EtcdServer) pb.KVServer { 52 return "aKVServer{ 53 NewKVServer(s), 54 quotaAlarmer{etcdserver.NewBackendQuota(s), s, s.ID()}, 55 } 56} 57 58func (s *quotaKVServer) Put(ctx context.Context, r *pb.PutRequest) (*pb.PutResponse, error) { 59 if err := s.qa.check(ctx, r); err != nil { 60 return nil, err 61 } 62 return s.KVServer.Put(ctx, r) 63} 64 65func (s *quotaKVServer) Txn(ctx context.Context, r *pb.TxnRequest) (*pb.TxnResponse, error) { 66 if err := s.qa.check(ctx, r); err != nil { 67 return nil, err 68 } 69 return s.KVServer.Txn(ctx, r) 70} 71 72type quotaLeaseServer struct { 73 pb.LeaseServer 74 qa quotaAlarmer 75} 76 77func (s *quotaLeaseServer) LeaseGrant(ctx context.Context, cr *pb.LeaseGrantRequest) (*pb.LeaseGrantResponse, error) { 78 if err := s.qa.check(ctx, cr); err != nil { 79 return nil, err 80 } 81 return s.LeaseServer.LeaseGrant(ctx, cr) 82} 83 84func NewQuotaLeaseServer(s *etcdserver.EtcdServer) pb.LeaseServer { 85 return "aLeaseServer{ 86 NewLeaseServer(s), 87 quotaAlarmer{etcdserver.NewBackendQuota(s), s, s.ID()}, 88 } 89} 90