|
@@ -0,0 +1,152 @@
|
|
|
+from multiprocessing import cpu_count, pool
|
|
|
+from os import getenv
|
|
|
+from random import randint, sample
|
|
|
+
|
|
|
+from asyncpg import create_pool
|
|
|
+from heaven import Application, Context, Request, Response
|
|
|
+from heaven.constants import STARTUP, SHUTDOWN
|
|
|
+from orjson import dumps
|
|
|
+
|
|
|
+
|
|
|
+#############
|
|
|
+# constants #
|
|
|
+#############
|
|
|
+APPLICATION_JSON = 'application/json'
|
|
|
+CONTENT_TYPE = 'Content-Type'
|
|
|
+POOL = 'pool'
|
|
|
+READ_ROW_SQL = 'SELECT "id", "randomnumber" FROM "world" WHERE id = $1'
|
|
|
+WRITE_ROW_SQL = 'UPDATE "world" SET "randomnumber"=$1 WHERE id=$2'
|
|
|
+ADDITIONAL_ROW = [0, "Additional fortune added at request time."]
|
|
|
+MAX_POOL_SIZE = 1000//cpu_count()
|
|
|
+MIN_POOL_SIZE = max(int(MAX_POOL_SIZE / 2), 1)
|
|
|
+
|
|
|
+
|
|
|
+#########################################
|
|
|
+# HOOKS for the app (shutdown, startup) #
|
|
|
+#########################################
|
|
|
+async def up_database(app: Application):
|
|
|
+ pool = await create_pool(
|
|
|
+ user=getenv("PGUSER", "benchmarkdbuser"),
|
|
|
+ password=getenv("PGPASS", "benchmarkdbpass"),
|
|
|
+ database="hello_world",
|
|
|
+ host="tfb-database",
|
|
|
+ port=5432,
|
|
|
+ min_size=MIN_POOL_SIZE,
|
|
|
+ max_size=MAX_POOL_SIZE,
|
|
|
+ )
|
|
|
+ app.keep(POOL, pool)
|
|
|
+
|
|
|
+
|
|
|
+async def down_database(app: Application):
|
|
|
+ pool = app.unkeep(POOL)
|
|
|
+ await pool.close()
|
|
|
+
|
|
|
+
|
|
|
+################
|
|
|
+# Helper utils #
|
|
|
+################
|
|
|
+def get_num_queries(queries):
|
|
|
+ try: query_count = int(queries)
|
|
|
+ except (ValueError, TypeError): return 1
|
|
|
+
|
|
|
+ if query_count < 1: return 1
|
|
|
+ if query_count > 500: return 500
|
|
|
+ return query_count
|
|
|
+
|
|
|
+
|
|
|
+###############################
|
|
|
+# Handlers for the app routes #
|
|
|
+###############################
|
|
|
+async def database(req: Request, res: Response, ctx: Context):
|
|
|
+ row_id = randint(1, 10000)
|
|
|
+ pool = req.app.peek(POOL)
|
|
|
+ async with pool.acquire() as connection:
|
|
|
+ number = await connection.fetchval(READ_ROW_SQL, row_id)
|
|
|
+
|
|
|
+ res.headers = CONTENT_TYPE, APPLICATION_JSON
|
|
|
+ res.body = dumps({"id": row_id, "randomNumber": number})
|
|
|
+
|
|
|
+
|
|
|
+async def json(req: Request, res: Response, ctx: Context):
|
|
|
+ res.headers = CONTENT_TYPE, APPLICATION_JSON
|
|
|
+ res.body = dumps({'message': 'Hello, World!'})
|
|
|
+
|
|
|
+
|
|
|
+async def queries(req: Request, res: Response, ctx: Context):
|
|
|
+ pool = req.app.peek(POOL)
|
|
|
+ queries = req.params.get('queries')
|
|
|
+ num_queries = get_num_queries(queries)
|
|
|
+ row_ids = sample(range(1, 10001), num_queries)
|
|
|
+ worlds = []
|
|
|
+
|
|
|
+ async with pool.acquire() as connection:
|
|
|
+ statement = await connection.prepare(READ_ROW_SQL)
|
|
|
+ for row_id in row_ids:
|
|
|
+ number = await statement.fetchval(row_id)
|
|
|
+ worlds.append({"id": row_id, "randomNumber": number})
|
|
|
+
|
|
|
+ res.headers = CONTENT_TYPE, APPLICATION_JSON
|
|
|
+ res.body = dumps(worlds)
|
|
|
+
|
|
|
+
|
|
|
+async def fortunes(req: Request, res: Response, ctx: Context):
|
|
|
+ pool = req.app.peek(POOL)
|
|
|
+ async with pool.acquire() as connection:
|
|
|
+ fortunes = await connection.fetch("SELECT * FROM Fortune")
|
|
|
+
|
|
|
+ fortunes.append(ADDITIONAL_ROW)
|
|
|
+ fortunes.sort(key=lambda row: row[1])
|
|
|
+ await res.render("fortune.html", fortunes=fortunes, request=req)
|
|
|
+
|
|
|
+
|
|
|
+async def updates(req: Request, res: Response, ctx: Context):
|
|
|
+ pool = req.app.peek(POOL)
|
|
|
+ queries = req.params.get('queries')
|
|
|
+ num_queries = get_num_queries(queries)
|
|
|
+
|
|
|
+ # To avoid deadlock
|
|
|
+ ids = sorted(sample(range(1, 10000 + 1), num_queries))
|
|
|
+ numbers = sorted(sample(range(1, 10000), num_queries))
|
|
|
+ updates = list(zip(ids, numbers))
|
|
|
+
|
|
|
+ worlds = [
|
|
|
+ {"id": row_id, "randomNumber": number} for row_id, number in updates
|
|
|
+ ]
|
|
|
+
|
|
|
+ async with pool.acquire() as connection:
|
|
|
+ statement = await connection.prepare(READ_ROW_SQL)
|
|
|
+ for row_id, _ in updates:
|
|
|
+ await statement.fetchval(row_id)
|
|
|
+ await connection.executemany(WRITE_ROW_SQL, updates)
|
|
|
+ res.headers = CONTENT_TYPE, APPLICATION_JSON
|
|
|
+ res.body = dumps(worlds)
|
|
|
+
|
|
|
+
|
|
|
+async def plaintext(req: Request, res: Response, ctx: Context):
|
|
|
+ res.headers = 'Content-Type', 'text/plain'
|
|
|
+ res.body = b"Hello, World!"
|
|
|
+
|
|
|
+
|
|
|
+################
|
|
|
+# App creation #
|
|
|
+################
|
|
|
+app = Application()
|
|
|
+app.TEMPLATES('templates')
|
|
|
+
|
|
|
+
|
|
|
+###################
|
|
|
+# Register hooks #
|
|
|
+###################
|
|
|
+app.ON(STARTUP, up_database)
|
|
|
+app.ON(SHUTDOWN, down_database)
|
|
|
+
|
|
|
+
|
|
|
+###################
|
|
|
+# Register routes #
|
|
|
+###################
|
|
|
+app.GET('/db', database)
|
|
|
+app.GET('/queries', queries)
|
|
|
+app.GET('/fortunes', fortunes)
|
|
|
+app.GET('/updates', updates)
|
|
|
+app.GET('/plaintext', plaintext)
|
|
|
+app.GET('/json', json)
|