WIP: multiprocessed bot, added avoidance
This commit is contained in:
parent
07ff2dcf01
commit
6cce15b2ca
GPG Key ID:
F0F6CFE11CDB0960
|
|
@ -0,0 +1,167 @@
|
|||
from multiprocessing.shared_memory import SharedMemory
|
||||
from typing import Optional
|
||||
import asyncio
|
||||
import socketio
|
||||
import aiohttp
|
||||
from aiohttp_socks import ProxyConnector
|
||||
from PIL import (
|
||||
Image,
|
||||
ImageFont,
|
||||
ImageDraw,
|
||||
ImageFilter,
|
||||
ImageSequence,
|
||||
ImageChops,
|
||||
)
|
||||
from enum import IntFlag
|
||||
from base64 import b64decode
|
||||
import signal
|
||||
import os
|
||||
import time
|
||||
|
||||
|
||||
class PixelMask(IntFlag):
|
||||
AVOID = 16
|
||||
MASK = 32
|
||||
FILL = 64
|
||||
CHECKED = 128
|
||||
|
||||
|
||||
class Manager:
|
||||
def __init__(self):
|
||||
self.shmem: Optional[SharedMemory] = None
|
||||
self.shmem_name = "omcb-bot"
|
||||
self.base = "https://onemillioncheckboxes.com"
|
||||
self.last_update = 0
|
||||
|
||||
self.bits_toggled_on = 0
|
||||
self.bits_toggled_off = 0
|
||||
self.last_printout = 0
|
||||
|
||||
async def listener(self):
|
||||
sio = socketio.AsyncClient()
|
||||
sio.on("connect", self.on_connect)
|
||||
sio.on("batched_bit_toggles", self.on_batched_bit_toggles)
|
||||
sio.on("full_state", self.on_full_state)
|
||||
|
||||
await sio.connect(f"{self.base}")
|
||||
await sio.wait()
|
||||
|
||||
def update_shmem(self, state: bytes):
|
||||
if not self.shmem:
|
||||
raise ValueError("shared memory is not initialized yet")
|
||||
|
||||
buffer = bytearray(bytes(1000000))
|
||||
for i in range(1000000):
|
||||
byte, bit = divmod(i, 8)
|
||||
if state[byte] & (0x80 >> bit):
|
||||
buffer[i] |= PixelMask.CHECKED
|
||||
else:
|
||||
buffer[i] &= ~PixelMask.CHECKED
|
||||
self.shmem.buf[:] = buffer
|
||||
|
||||
async def on_connect(self):
|
||||
async with aiohttp.ClientSession() as http:
|
||||
async with http.get(f"{self.base}/api/initial-state") as req:
|
||||
data = await req.json()
|
||||
buffer = b64decode(data["full_state"].encode() + b"=")
|
||||
self.update_shmem(buffer)
|
||||
self.last_update = data["timestamp"]
|
||||
|
||||
async def on_full_state(self, data):
|
||||
if not self.shmem:
|
||||
raise ValueError("shared memory is not initialized yet")
|
||||
buffer = b64decode(data["full_state"].encode() + b"=")
|
||||
self.update_shmem(buffer)
|
||||
self.last_update = data["timestamp"]
|
||||
|
||||
async def on_batched_bit_toggles(self, data):
|
||||
if not self.shmem:
|
||||
raise ValueError("shared memory is not initialized yet")
|
||||
bits_on, bits_off, timestamp = data
|
||||
if timestamp < self.last_update:
|
||||
print("old update, ignoring")
|
||||
|
||||
self.last_update = timestamp
|
||||
|
||||
self.bits_toggled_on += len(bits_on)
|
||||
self.bits_toggled_off = len(bits_off)
|
||||
|
||||
for bit in bits_on:
|
||||
self.shmem.buf[bit] |= PixelMask.CHECKED
|
||||
for bit in bits_off:
|
||||
self.shmem.buf[bit] &= ~PixelMask.CHECKED
|
||||
|
||||
since_last_printout = time.time() - self.last_printout
|
||||
if since_last_printout >= 5:
|
||||
self.last_printout = time.time()
|
||||
print()
|
||||
print(f"Toggled on: {self.bits_toggled_on / since_last_printout}/s")
|
||||
print(f"Toggled off: {self.bits_toggled_off / since_last_printout}/s")
|
||||
self.bits_toggled_on = self.bits_toggled_off = 0
|
||||
|
||||
def on_sigusr1(self, signum, frame):
|
||||
if not self.shmem:
|
||||
raise ValueError("shared memory is not initialized yet")
|
||||
print("Caught SIGUSR1, dumping state")
|
||||
buf = bytes(self.shmem.buf[:])
|
||||
with Image.new("RGB", (1000, 1000), 0) as im:
|
||||
for i in range(1000000):
|
||||
y, x = divmod(i, 1000)
|
||||
im.putpixel((x, y), (
|
||||
255 if buf[i] & PixelMask.FILL else 0,
|
||||
255 if buf[i] & PixelMask.MASK else 0,
|
||||
255 if buf[i] & PixelMask.CHECKED else 0
|
||||
))
|
||||
im.save("state.png")
|
||||
|
||||
async def animator(self):
|
||||
while True:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
def add_avoid_range(self, rng: range):
|
||||
if not self.shmem:
|
||||
raise ValueError("shared memory is not initialized yet")
|
||||
for ndx in rng:
|
||||
self.shmem.buf[ndx] |= PixelMask.AVOID
|
||||
|
||||
def add_avoid_rect(self, sx: int, sy: int, w: int, h: int):
|
||||
for y in range(sy, sy + h):
|
||||
ox = y * 1000
|
||||
self.add_avoid_range(range(sx + ox, sx + w + ox))
|
||||
|
||||
def add_avoid_index(self, index: int):
|
||||
if not self.shmem:
|
||||
raise ValueError("shared memory is not initialized yet")
|
||||
assert 0 <= index < 1000000
|
||||
self.shmem.buf[index] |= PixelMask.AVOID
|
||||
|
||||
def add_avoid_image(self, im: Image.Image):
|
||||
assert im.width == 1000
|
||||
assert im.height == 1000
|
||||
for i in range(1000000):
|
||||
y, x = divmod(i, 1000)
|
||||
la = im.getpixel((x, y))
|
||||
assert isinstance(la, (tuple, list)) and len(la) == 2
|
||||
if la[1]:
|
||||
self.add_avoid_index(i)
|
||||
|
||||
async def __aenter__(self):
|
||||
self.shmem = SharedMemory(self.shmem_name, True, 1000000)
|
||||
return self
|
||||
|
||||
async def __aexit__(self, a, b, c):
|
||||
if self.shmem:
|
||||
print("cleaning up shmem")
|
||||
self.shmem.close()
|
||||
self.shmem.unlink()
|
||||
|
||||
|
||||
async def main():
|
||||
print(f"PID: {os.getpid()}")
|
||||
async with Manager() as mgr:
|
||||
signal.signal(signal.SIGUSR1, mgr.on_sigusr1)
|
||||
await mgr.listener()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Loading…
Reference in New Issue