code.grnet.gr Git - aquarium/blob - src/main/scala/gr/grnet/aquarium/store/mongodb/MongoDBStore.scala

   1 /*
   2  * Copyright 2011-2012 GRNET S.A. All rights reserved.
   3  *
   4  * Redistribution and use in source and binary forms, with or
   5  * without modification, are permitted provided that the following
   6  * conditions are met:
   7  *
   8  *   1. Redistributions of source code must retain the above
   9  *      copyright notice, this list of conditions and the following
  10  *      disclaimer.
  11  *
  12  *   2. Redistributions in binary form must reproduce the above
  13  *      copyright notice, this list of conditions and the following
  14  *      disclaimer in the documentation and/or other materials
  15  *      provided with the distribution.
  16  *
  17  * THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS
  18  * OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  19  * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
  20  * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR
  21  * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  22  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  23  * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
  24  * USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
  25  * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  26  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
  27  * ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  28  * POSSIBILITY OF SUCH DAMAGE.
  29  *
  30  * The views and conclusions contained in the software and
  31  * documentation are those of the authors and should not be
  32  * interpreted as representing official policies, either expressed
  33  * or implied, of GRNET S.A.
  34  */
  35
  36 package gr.grnet.aquarium.store.mongodb
  37
  38 import collection.immutable
  39 import com.mongodb._
  40 import gr.grnet.aquarium.computation.BillingMonthInfo
  41 import gr.grnet.aquarium.converter.StdConverters
  42 import gr.grnet.aquarium.logic.accounting.dsl.Timeslot
  43 import gr.grnet.aquarium.message.MessageConstants
  44 import gr.grnet.aquarium.message.avro.gen.{UserStateMsg, IMEventMsg, ResourceEventMsg, PolicyMsg}
  45 import gr.grnet.aquarium.message.avro.{MessageFactory, OrderingHelpers, AvroHelpers}
  46 import gr.grnet.aquarium.store._
  47 import gr.grnet.aquarium.util._
  48 import gr.grnet.aquarium.util.json.JsonSupport
  49 import gr.grnet.aquarium.{Aquarium, AquariumException}
  50 import org.apache.avro.specific.SpecificRecord
  51 import org.bson.types.ObjectId
  52
  53 /**
  54  * Mongodb implementation of the various aquarium stores.
  55  *
  56  * @author Christos KK Loverdos <loverdos@gmail.com>
  57  * @author Georgios Gousios <gousiosg@gmail.com>
  58  * @author Prodromos Gerakios <pgerakio@grnet.gr>
  59  */
  60 class MongoDBStore(
  61     val aquarium: Aquarium,
  62     val mongo: Mongo,
  63     val database: String,
  64     val username: String,
  65     val password: String)
  66   extends ResourceEventStore
  67   with UserStateStore
  68   with IMEventStore
  69   with PolicyStore
  70   with Loggable {
  71
  72   private[store] lazy val resourceEvents = getCollection(MongoDBStore.ResourceEventCollection)
  73   private[store] lazy val userStates = getCollection(MongoDBStore.UserStateCollection)
  74   private[store] lazy val imEvents = getCollection(MongoDBStore.IMEventCollection)
  75   private[store] lazy val policies = getCollection(MongoDBStore.PolicyCollection)
  76
  77   private[this] def getCollection(name: String): DBCollection = {
  78     val db = mongo.getDB(database)
  79     //logger.debug("Authenticating to mongo")
  80     if(!db.isAuthenticated && !db.authenticate(username, password.toCharArray)) {
  81       throw new AquariumException("Could not authenticate user %s".format(username))
  82     }
  83     db.getCollection(name)
  84   }
  85
  86   //+ResourceEventStore
  87   def pingResourceEventStore(): Unit = synchronized {
  88     MongoDBStore.ping(mongo)
  89   }
  90
  91   def insertResourceEvent(event: ResourceEventMsg) = {
  92     val mongoID = new ObjectId()
  93     event.setInStoreID(mongoID.toStringMongod)
  94
  95     val dbObject = new BasicDBObjectBuilder().
  96       add(MongoDBStore.JsonNames._id, mongoID).
  97       add(MongoDBStore.JsonNames.payload, AvroHelpers.bytesOfSpecificRecord(event)).
  98       add(MongoDBStore.JsonNames.userID, event.getUserID).
  99       add(MongoDBStore.JsonNames.occurredMillis, event.getOccurredMillis).
 100       add(MongoDBStore.JsonNames.receivedMillis, event.getReceivedMillis).
 101     get()
 102
 103     MongoDBStore.insertDBObject(dbObject, resourceEvents)
 104     event
 105   }
 106
 107   def findResourceEventByID(id: String): Option[ResourceEventMsg] = {
 108     val dbObjectOpt = MongoDBStore.findOneByAttribute(resourceEvents, MongoDBStore.JsonNames.id, id)
 109     for {
 110       dbObject ← dbObjectOpt
 111       payload = dbObject.get(MongoDBStore.JsonNames.payload)
 112       msg = AvroHelpers.specificRecordOfBytes(payload.asInstanceOf[Array[Byte]], new ResourceEventMsg)
 113     } yield msg
 114   }
 115
 116   def countOutOfSyncResourceEventsForBillingPeriod(userID: String, startMillis: Long, stopMillis: Long): Long = {
 117     val query = new BasicDBObjectBuilder().
 118       add(MongoDBStore.JsonNames.userID, userID).
 119       // received within the period
 120       add(MongoDBStore.JsonNames.receivedMillis, new BasicDBObject("$gte", startMillis)).
 121       add(MongoDBStore.JsonNames.receivedMillis, new BasicDBObject("$lte", stopMillis)).
 122       // occurred outside the period
 123       add("$or", {
 124         val dbList = new BasicDBList()
 125         dbList.add(0, new BasicDBObject(MongoDBStore.JsonNames.occurredMillis, new BasicDBObject("$lt", startMillis)))
 126         dbList.add(1, new BasicDBObject(MongoDBStore.JsonNames.occurredMillis, new BasicDBObject("$gt", stopMillis)))
 127         dbList
 128       }).
 129       get()
 130
 131     resourceEvents.count(query)
 132   }
 133
 134   def foreachResourceEventOccurredInPeriod(
 135       userID: String,
 136       startMillis: Long,
 137       stopMillis: Long
 138   )(f: ResourceEventMsg ⇒ Unit): Unit = {
 139
 140     val query = new BasicDBObjectBuilder().
 141       add(MongoDBStore.JsonNames.userID, userID).
 142       add(MongoDBStore.JsonNames.occurredMillis, new BasicDBObject("$gte", startMillis)).
 143       add(MongoDBStore.JsonNames.occurredMillis, new BasicDBObject("$lte", stopMillis)).
 144       get()
 145
 146     val sorter = new BasicDBObject(MongoDBStore.JsonNames.occurredMillis, 1)
 147     val cursor = resourceEvents.find(query).sort(sorter)
 148
 149     withCloseable(cursor) { cursor ⇒
 150       while(cursor.hasNext) {
 151         val nextDBObject = cursor.next()
 152         val payload = nextDBObject.get(MongoDBStore.JsonNames.payload).asInstanceOf[Array[Byte]]
 153         val nextEvent = AvroHelpers.specificRecordOfBytes(payload, new ResourceEventMsg)
 154
 155         f(nextEvent)
 156       }
 157     }
 158   }
 159   //-ResourceEventStore
 160
 161   //+ UserStateStore
 162   def findUserStateByUserID(userID: String) = {
 163     val dbObjectOpt = MongoDBStore.findOneByAttribute(userStates, MongoDBStore.JsonNames.userID, userID)
 164     for {
 165       dbObject <- dbObjectOpt
 166       payload = dbObject.get(MongoDBStore.JsonNames.payload).asInstanceOf[Array[Byte]]
 167       msg = AvroHelpers.specificRecordOfBytes(payload, new UserStateMsg)
 168     } yield {
 169       msg
 170     }
 171   }
 172
 173   def findLatestUserStateForFullMonthBilling(userID: String, bmi: BillingMonthInfo) = {
 174     val query = new BasicDBObjectBuilder().
 175       add(MongoDBStore.JsonNames.userID, userID).
 176       add(MongoDBStore.JsonNames.isFullBillingMonth, true).
 177       add(MongoDBStore.JsonNames.billingYear, bmi.year).
 178       add(MongoDBStore.JsonNames.billingMonth, bmi.month).
 179       get()
 180
 181     // Descending order, so that the latest comes first
 182     val sorter = new BasicDBObject(MongoDBStore.JsonNames.occurredMillis, -1)
 183
 184     val cursor = userStates.find(query).sort(sorter)
 185
 186     withCloseable(cursor) { cursor ⇒
 187       MongoDBStore.findNextPayloadRecord(cursor, new UserStateMsg)
 188     }
 189   }
 190
 191   /**
 192    * Stores a user state.
 193    */
 194   def insertUserState(event: UserStateMsg)= {
 195     val mongoID = new ObjectId()
 196     event.setInStoreID(mongoID.toStringMongod)
 197
 198     val dbObject = new BasicDBObjectBuilder().
 199       add(MongoDBStore.JsonNames._id, mongoID).
 200       add(MongoDBStore.JsonNames.payload, AvroHelpers.bytesOfSpecificRecord(event)).
 201       add(MongoDBStore.JsonNames.userID, event.getUserID).
 202       add(MongoDBStore.JsonNames.occurredMillis, event.getOccurredMillis).
 203       add(MongoDBStore.JsonNames.isFullBillingMonth, event.getIsFullBillingMonth).
 204       add(MongoDBStore.JsonNames.billingYear, event.getBillingYear).
 205       add(MongoDBStore.JsonNames.billingMonth, event.getBillingMonth).
 206       add(MongoDBStore.JsonNames.billingMonthDay, event.getBillingMonthDay).
 207     get()
 208
 209     MongoDBStore.insertDBObject(dbObject, userStates)
 210     event
 211   }
 212   //- UserStateStore
 213
 214   //+IMEventStore
 215   def pingIMEventStore(): Unit = {
 216     MongoDBStore.ping(mongo)
 217   }
 218
 219   def insertIMEvent(event: IMEventMsg) = {
 220     val mongoID = new ObjectId()
 221     event.setInStoreID(mongoID.toStringMongod)
 222
 223     val dbObject = new BasicDBObjectBuilder().
 224       add(MongoDBStore.JsonNames._id, mongoID).
 225       add(MongoDBStore.JsonNames.payload, AvroHelpers.bytesOfSpecificRecord(event)).
 226       add(MongoDBStore.JsonNames.userID, event.getUserID).
 227       add(MongoDBStore.JsonNames.eventType, event.getEventType().toLowerCase).
 228       add(MongoDBStore.JsonNames.occurredMillis, event.getOccurredMillis).
 229       add(MongoDBStore.JsonNames.receivedMillis, event.getReceivedMillis).
 230     get()
 231
 232     MongoDBStore.insertDBObject(dbObject, imEvents)
 233     event
 234   }
 235
 236   def findIMEventByID(id: String) = {
 237     val dbObjectOpt = MongoDBStore.findOneByAttribute(imEvents, MongoDBStore.JsonNames.id, id)
 238     for {
 239       dbObject ← dbObjectOpt
 240       payload = dbObject.get(MongoDBStore.JsonNames.payload).asInstanceOf[Array[Byte]]
 241       msg = AvroHelpers.specificRecordOfBytes(payload, new IMEventMsg)
 242     } yield {
 243       msg
 244     }
 245   }
 246
 247
 248   /**
 249    * Find the `CREATE` even for the given user. Note that there must be only one such event.
 250    */
 251   def findCreateIMEventByUserID(userID: String) = {
 252     val query = new BasicDBObjectBuilder().
 253       add(MongoDBStore.JsonNames.userID, userID).
 254       add(MongoDBStore.JsonNames.eventType, MessageConstants.IMEventMsg.EventTypes.create).get()
 255
 256     // Normally one such event is allowed ...
 257     val cursor = imEvents.find(query).sort(new BasicDBObject(MongoDBStore.JsonNames.occurredMillis, 1))
 258
 259     val dbObjectOpt = withCloseable(cursor) { cursor ⇒
 260       if(cursor.hasNext) {
 261         Some(cursor.next())
 262       } else {
 263         None
 264       }
 265     }
 266
 267     for {
 268       dbObject <- dbObjectOpt
 269       payload = dbObject.get(MongoDBStore.JsonNames.payload).asInstanceOf[Array[Byte]]
 270       msg = AvroHelpers.specificRecordOfBytes(payload, new IMEventMsg)
 271     } yield {
 272       msg
 273     }
 274   }
 275
 276   /**
 277    * Scans events for the given user, sorted by `occurredMillis` in ascending order and runs them through
 278    * the given function `f`.
 279    *
 280    * Any exception is propagated to the caller. The underlying DB resources are properly disposed in any case.
 281    */
 282   def foreachIMEventInOccurrenceOrder(userID: String)(f: (IMEventMsg) ⇒ Unit) = {
 283     val query = new BasicDBObject(MongoDBStore.JsonNames.userID, userID)
 284     val cursor = imEvents.find(query).sort(new BasicDBObject(MongoDBStore.JsonNames.occurredMillis, 1))
 285
 286     withCloseable(cursor) { cursor ⇒
 287       while(cursor.hasNext) {
 288         val dbObject = cursor.next()
 289         val payload = dbObject.get(MongoDBStore.JsonNames.payload).asInstanceOf[Array[Byte]]
 290         val msg = AvroHelpers.specificRecordOfBytes(payload, new IMEventMsg)
 291
 292         f(msg)
 293       }
 294     }
 295   }
 296   //-IMEventStore
 297
 298
 299
 300   //+PolicyStore
 301   def foreachPolicy[U](f: PolicyMsg ⇒ U) {
 302     val cursor = policies.find()
 303     withCloseable(cursor) { cursor ⇒
 304       while(cursor.hasNext) {
 305         val dbObject = cursor.next()
 306         val payload = dbObject.get(MongoDBStore.JsonNames.payload).asInstanceOf[Array[Byte]]
 307         val policy = AvroHelpers.specificRecordOfBytes(payload, new PolicyMsg)
 308         f(policy)
 309       }
 310     }
 311   }
 312
 313   def insertPolicy(policy: PolicyMsg): PolicyMsg = {
 314     val mongoID = new ObjectId()
 315     policy.setInStoreID(mongoID.toStringMongod)
 316     val dbObject = new BasicDBObjectBuilder().
 317       add(MongoDBStore.JsonNames._id, mongoID).
 318       add(MongoDBStore.JsonNames.validFromMillis, policy.getValidFromMillis).
 319       add(MongoDBStore.JsonNames.validToMillis, policy.getValidToMillis).
 320       add(MongoDBStore.JsonNames.payload, AvroHelpers.bytesOfSpecificRecord(policy)).
 321     get()
 322
 323     MongoDBStore.insertDBObject(dbObject, policies)
 324     policy
 325   }
 326
 327   def loadPolicyAt(atMillis: Long): Option[PolicyMsg] = {
 328     // FIXME Inefficient
 329     var _policies = immutable.TreeSet[PolicyMsg]()(OrderingHelpers.DefaultPolicyMsgOrdering)
 330     foreachPolicy(_policies += _)
 331     _policies.to(MessageFactory.newDummyPolicyMsgAt(atMillis)).lastOption
 332   }
 333
 334   def loadSortedPoliciesWithin(fromMillis: Long, toMillis: Long): immutable.SortedMap[Timeslot, PolicyMsg] = {
 335     // FIXME Inefficient
 336     var _policies = immutable.TreeSet[PolicyMsg]()(OrderingHelpers.DefaultPolicyMsgOrdering)
 337     foreachPolicy(_policies += _)
 338
 339     immutable.SortedMap(_policies.
 340       from(MessageFactory.newDummyPolicyMsgAt(fromMillis)).
 341       to(MessageFactory.newDummyPolicyMsgAt(toMillis)).toSeq.
 342       map(p ⇒ (Timeslot(p.getValidFromMillis, p.getValidToMillis), p)): _*
 343     )
 344   }
 345   //-PolicyStore
 346 }
 347
 348 object MongoDBStore {
 349   final val JsonNames = gr.grnet.aquarium.util.json.JsonNames
 350
 351   final val ResourceEventCollection = "resevents"
 352
 353   final val UserStateCollection = "userstates"
 354
 355   final val IMEventCollection = "imevents"
 356
 357   final val PolicyCollection = "policies"
 358
 359   def firstResultIfExists[A](cursor: DBCursor, f: DBObject ⇒ A): Option[A] = {
 360     withCloseable(cursor) { cursor ⇒
 361       if(cursor.hasNext) {
 362         Some(f(cursor.next()))
 363       } else {
 364         None
 365       }
 366     }
 367   }
 368
 369   def ping(mongo: Mongo): Unit = synchronized {
 370     // This requires a network roundtrip
 371     mongo.isLocked
 372   }
 373
 374   def findOneByAttribute(
 375       collection: DBCollection,
 376       attributeName: String,
 377       attributeValue: String,
 378       sortByOpt: Option[DBObject] = None
 379   ): Option[DBObject] =  {
 380     val query = new BasicDBObject(attributeName, attributeValue)
 381     val cursor = sortByOpt match {
 382       case None         ⇒ collection find query
 383       case Some(sortBy) ⇒ collection find query sort sortBy
 384     }
 385     withCloseable(cursor) { cursor ⇒
 386       if(cursor.hasNext) Some(cursor.next()) else None
 387     }
 388   }
 389
 390   def insertDBObject(dbObj: DBObject, collection: DBCollection) {
 391     collection.insert(dbObj, WriteConcern.JOURNAL_SAFE)
 392   }
 393
 394   def findNextPayloadRecord[R <: SpecificRecord](cursor: DBCursor, fresh: R): Option[R] = {
 395     for {
 396       dbObject <- if(cursor.hasNext) Some(cursor.next()) else None
 397       payload = dbObject.get(MongoDBStore.JsonNames.payload).asInstanceOf[Array[Byte]]
 398       msg = AvroHelpers.specificRecordOfBytes(payload, fresh)
 399     } yield {
 400       msg
 401     }
 402   }
 403
 404   def jsonSupportToDBObject(jsonSupport: JsonSupport) = {
 405     StdConverters.AllConverters.convertEx[DBObject](jsonSupport)
 406   }
 407 }