@@ -17,15 +17,17 @@ import (
1717
1818// BoltDB is a database using BoltDB
1919type BoltDB struct {
20- db * bolt.DB
21- locksBucketName []byte
22- pullsBucketName []byte
20+ db * bolt.DB
21+ locksBucketName []byte
22+ pullsBucketName []byte
23+ globalLocksBucketName []byte
2324}
2425
2526const (
26- locksBucketName = "runLocks"
27- pullsBucketName = "pulls"
28- pullKeySeparator = "::"
27+ locksBucketName = "runLocks"
28+ pullsBucketName = "pulls"
29+ globalLocksBucketName = "globalLocks"
30+ pullKeySeparator = "::"
2931)
3032
3133// New returns a valid locker. We need to be able to write to dataDir
@@ -50,18 +52,31 @@ func New(dataDir string) (*BoltDB, error) {
5052 if _ , err = tx .CreateBucketIfNotExists ([]byte (pullsBucketName )); err != nil {
5153 return errors .Wrapf (err , "creating bucket %q" , pullsBucketName )
5254 }
55+ if _ , err = tx .CreateBucketIfNotExists ([]byte (globalLocksBucketName )); err != nil {
56+ return errors .Wrapf (err , "creating bucket %q" , globalLocksBucketName )
57+ }
5358 return nil
5459 })
5560 if err != nil {
5661 return nil , errors .Wrap (err , "starting BoltDB" )
5762 }
5863 // todo: close BoltDB when server is sigtermed
59- return & BoltDB {db : db , locksBucketName : []byte (locksBucketName ), pullsBucketName : []byte (pullsBucketName )}, nil
64+ return & BoltDB {
65+ db : db ,
66+ locksBucketName : []byte (locksBucketName ),
67+ pullsBucketName : []byte (pullsBucketName ),
68+ globalLocksBucketName : []byte (globalLocksBucketName ),
69+ }, nil
6070}
6171
6272// NewWithDB is used for testing.
63- func NewWithDB (db * bolt.DB , bucket string ) (* BoltDB , error ) {
64- return & BoltDB {db : db , locksBucketName : []byte (bucket ), pullsBucketName : []byte (pullsBucketName )}, nil
73+ func NewWithDB (db * bolt.DB , bucket string , globalBucket string ) (* BoltDB , error ) {
74+ return & BoltDB {
75+ db : db ,
76+ locksBucketName : []byte (bucket ),
77+ pullsBucketName : []byte (pullsBucketName ),
78+ globalLocksBucketName : []byte (globalBucket ),
79+ }, nil
6580}
6681
6782// TryLock attempts to create a new lock. If the lock is
@@ -155,6 +170,87 @@ func (b *BoltDB) List() ([]models.ProjectLock, error) {
155170 return locks , nil
156171}
157172
173+ // LockCommand attempts to create a new lock for a CommandName.
174+ // If the lock doesn't exists, it will create a lock and return a pointer to it.
175+ // If the lock already exists, it will return an "lock already exists" error
176+ func (b * BoltDB ) LockCommand (cmdName models.CommandName , lockTime time.Time ) (* models.CommandLock , error ) {
177+ lock := models.CommandLock {
178+ CommandName : cmdName ,
179+ LockMetadata : models.LockMetadata {
180+ UnixTime : lockTime .Unix (),
181+ },
182+ }
183+
184+ newLockSerialized , _ := json .Marshal (lock )
185+ transactionErr := b .db .Update (func (tx * bolt.Tx ) error {
186+ bucket := tx .Bucket (b .globalLocksBucketName )
187+
188+ currLockSerialized := bucket .Get ([]byte (b .commandLockKey (cmdName )))
189+ if currLockSerialized != nil {
190+ return errors .New ("lock already exists" )
191+ }
192+
193+ // This will only error on readonly buckets, it's okay to ignore.
194+ bucket .Put ([]byte (b .commandLockKey (cmdName )), newLockSerialized ) // nolint: errcheck
195+ return nil
196+ })
197+
198+ if transactionErr != nil {
199+ return nil , errors .Wrap (transactionErr , "db transaction failed" )
200+ }
201+
202+ return & lock , nil
203+ }
204+
205+ // UnlockCommand removes CommandName lock if present.
206+ // If there are no lock it returns an error.
207+ func (b * BoltDB ) UnlockCommand (cmdName models.CommandName ) error {
208+ transactionErr := b .db .Update (func (tx * bolt.Tx ) error {
209+ bucket := tx .Bucket (b .globalLocksBucketName )
210+
211+ if l := bucket .Get ([]byte (b .commandLockKey (cmdName ))); l == nil {
212+ return errors .New ("no lock exists" )
213+ }
214+
215+ return bucket .Delete ([]byte (b .commandLockKey (cmdName )))
216+ })
217+
218+ if transactionErr != nil {
219+ return errors .Wrap (transactionErr , "db transaction failed" )
220+ }
221+
222+ return nil
223+ }
224+
225+ // CheckCommandLock checks if CommandName lock was set.
226+ // If the lock exists return the pointer to the lock object, otherwise return nil
227+ func (b * BoltDB ) CheckCommandLock (cmdName models.CommandName ) (* models.CommandLock , error ) {
228+ cmdLock := models.CommandLock {}
229+
230+ found := false
231+
232+ err := b .db .View (func (tx * bolt.Tx ) error {
233+ bucket := tx .Bucket (b .globalLocksBucketName )
234+
235+ serializedLock := bucket .Get ([]byte (b .commandLockKey (cmdName )))
236+
237+ if serializedLock != nil {
238+ if err := json .Unmarshal (serializedLock , & cmdLock ); err != nil {
239+ return errors .Wrap (err , "failed to deserialize UserConfig" )
240+ }
241+ found = true
242+ }
243+
244+ return nil
245+ })
246+
247+ if found {
248+ return & cmdLock , err
249+ }
250+
251+ return nil , err
252+ }
253+
158254// UnlockByPull deletes all locks associated with that pull request and returns them.
159255func (b * BoltDB ) UnlockByPull (repoFullName string , pullNum int ) ([]models.ProjectLock , error ) {
160256 var locks []models.ProjectLock
@@ -355,6 +451,10 @@ func (b *BoltDB) pullKey(pull models.PullRequest) ([]byte, error) {
355451 nil
356452}
357453
454+ func (b * BoltDB ) commandLockKey (cmdName models.CommandName ) string {
455+ return fmt .Sprintf ("%s/lock" , cmdName )
456+ }
457+
358458func (b * BoltDB ) lockKey (p models.Project , workspace string ) string {
359459 return fmt .Sprintf ("%s/%s/%s" , p .RepoFullName , p .Path , workspace )
360460}
0 commit comments