command-handler.js

const { serialize, serializeError } = require("./protocol/serialize")
const {
  config,
  replicationState,
  replicas,
  addReplicaSocket,
  updateReplicaOffset,
} = require("./config/config")
const { toIPV4 } = require("./utils")
const { map, expiryMap } = require("./store")
const { Worker } = require("worker_threads")
const { readFromBacklog, isIncludedInBacklog } = require("./backlog")
const { connectToMaster } = require("./utils/replica.util")

const RESET = "\x1b[0m"
const RED = "\x1b[31m"
const GREEN = "\x1b[32m"
const BLUE = "\x1b[36m"

const OK = "+OK\r\n"
const QUEUED = "+QUEUED\r\n"

const COMMANDS = {
  PING: "PING",
  SET: "SET",
  GET: "GET",
  TTL: "TTL",
  EXISTS: "EXISTS",
  DEL: "DEL",
  INCR: "INCR",
  DECR: "DECR",
  INFO: "INFO",
  REPLCONF: "REPLCONF",
  REPLICAOF: "REPLICAOF",
  PSYNC: "PSYNC",
  CONFIG: "CONFIG",
  SAVE: "SAVE",
  MULTI: "MULTI",
  EXEC: "EXEC",
  DISCARD: "DISCARD",
}

function queueCommand(args, socket) {
  socket.transactionArray.push(args)
  return QUEUED
}

function commandHandler(args, socket) {
  const command = args[0]
  switch (command) {
    case COMMANDS.PING:
      return handlePing(args)
    case COMMANDS.SET:
      return handleSet(args, socket)
    case COMMANDS.GET:
      return handleGet(args, socket)
    case COMMANDS.TTL:
      return handleTtl(args, socket)
    case COMMANDS.EXISTS:
      return handleExists(args, socket)
    case COMMANDS.DEL:
      return handleDel(args, socket)
    case COMMANDS.INCR:
      return handleIncr(args, socket)
    case COMMANDS.DECR:
      return handleDecr(args, socket)
    case COMMANDS.INFO:
      return handleInfo()
    case COMMANDS.REPLCONF:
      return handleReplConf(args, socket)
    case COMMANDS.REPLICAOF:
      return handleReplicaOf(args)
    case COMMANDS.PSYNC:
      return handlePsync(args, socket)
    case COMMANDS.CONFIG:
      return handleConfig(args)
    case COMMANDS.SAVE:
      return handleSave(args)
    case COMMANDS.MULTI:
      return handleMulti(args, socket)
    case COMMANDS.EXEC:
      return handleExec(args, socket)
    case COMMANDS.DISCARD:
      return handleDiscard(args, socket)
    default:
      return "-Unknown command\r\n"
  }
}

function handlePing(args) {
  if (args.length === 1) return serialize("PONG")
  return serialize(args[1])
}

function handleSet(args, socket) {
  const length = args.length
  const [_, key, value, option, expireTime] = args

  if (length !== 3 && length !== 5) throw new Error("Invalid SET command")
  if (length === 5 && !["EX", "PX"].includes(option))
    throw new Error("Unknown Option specified")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  map.set(key, value)
  if (option === "EX") {
    expiryMap.set(key, Date.now() + expireTime * 1000)
  } else if (option === "PX") {
    expiryMap.set(key, Date.now() + Number(expireTime))
  }
  if (replicationState.isMaster) {
    return OK
  }
}

function handleGet(args, socket) {
  if (args.length !== 2) throw new Error("Invalid GET command")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  const key = args[1]
  const timestamp = expiryMap.get(key)
  if (timestamp && Date.now() >= timestamp) {
    expiryMap.delete(key)
    map.delete(key)
    return serializeError("Key expired")
  }
  const value = map.get(key)
  if (value === undefined) {
    return serializeError("Key doesn't exist")
  }
  return serialize(value)
}

function handleTtl(args, socket) {
  if (args.length !== 2) throw new Error("Invalid TTL command")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  const key = args[1]
  const timestamp = expiryMap.get(key)
  if (!timestamp) {
    return serializeError("Key doesn't exist")
  }
  if (Date.now() >= timestamp) {
    expiryMap.delete(key)
    map.delete(key)
    return serializeError("Key expired")
  }
  const remainingSeconds = Math.floor((timestamp - Date.now()) / 1000)

  return serialize(remainingSeconds)
}

function handleExists(args, socket) {
  if (args.length < 2) throw new Error("Invalid Exists command")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  const [_, ...keys] = args
  let count = 0
  keys.forEach((key) => {
    if (map.has(key)) {
      count++
    }
  })

  return serialize(count)
}

function handleDel(args, socket) {
  if (args.length < 2) throw new Error("Invalid Del command")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  const [_, ...keys] = args
  let count = 0
  keys.forEach((key) => {
    if (map.has(key)) {
      map.delete(key)
      count++
    }
  })

  if (replicationState.isMaster) {
    return serialize(count)
  }
}

function handleIncr(args, socket) {
  if (args.length !== 2) throw new Error("Invalid Incr command")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  const [_, key] = args
  let value = map.get(key)

  if (typeof value !== "number") {
    return serializeError("Key type is not an integer")
  }

  if (value === undefined) {
    map.set(key, 1)
    if (replicationState.isMaster) {
      return serialize(1)
    }
  }
  value++
  map.set(key, value)
  if (replicationState.isMaster) {
    return serialize(value)
  }
}

function handleDecr(args, socket) {
  if (args.length !== 2) throw new Error("Invalid Decr command")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  const [_, key] = args
  let value = map.get(key)

  if (value === undefined) {
    return serializeError("Key doesn't exist")
  }

  if (typeof value !== "number") {
    return serializeError("Key type is not an integer")
  }

  map.set(key, --value)
  if (replicationState.isMaster) {
    return serialize(value)
  }
}

function handleConfig(args, socket) {
  if (args.length !== 3) throw new Error("Invalid CONFIG command")

  if (replicationState.isMaster && socket.isMultiMode) {
    return queueCommand(args, socket)
  }

  const [_, configCommand, configKey] = args
  const key = configKey.toLowerCase()
  if (configCommand === "GET") {
    const value = config[key]
    if (!value) {
      return "-CONFIG KEY NOT FOUND\r\n"
    }
    return serialize([key, value])
  }
}

function handleInfo() {
  const mappedReplicas = replicas.sockets.map(
    ({ socket, port, offset }, i) =>
      `slave${i}:ip=${toIPV4(
        socket.localAddress
      )},port=${port},offset=${offset}`
  )

  return serialize(
    `# Replication\n` +
      `role:${
        replicationState.isMaster
          ? RED + "master" + RESET
          : GREEN + "slave" + RESET
      }\n` +
      `${
        replicationState.isMaster
          ? `connected_slaves:${replicas.sockets.length}\n`
          : ``
      }` +
      (mappedReplicas.length > 0 ? mappedReplicas.join("\n") + "\n" : "") +
      `replid:${replicationState.id}\n` +
      `repl_offset:${BLUE}${replicationState.offset}${RESET}`
  )
}

function handleReplConf(args, socket) {
  // Respond as replica
  if (args[1] === "GETACK" && args[2] === "*" && !replicationState.isMaster) {
    return serialize(["REPLCONF", "ACK", `${replicationState.offset}`])
  }
  if (args[1] === "listening-port" && replicationState.isMaster) {
    const port = args[2]
    addReplicaSocket(socket, port, -1)
    return "$10\r\nMASTER-ACK\r\n"
  }

  if (args[1] === "ACK" && replicationState.isMaster) {
    const replicaOffset = Number(args[2])
    const replica = replicas.sockets.find((r) => r.socket === socket)
    if (replica) {
      replica.offset = replicaOffset // update offset
    }

    return
  }

  return "$10\r\nMASTER-ACK\r\n"
}

function handleReplicaOf(args) {
  if (args.length !== 3) throw new Error("Invalid REPLICAOF command")

  const [_, host, port] = args
  if (host === "NO" && port === "ONE") {
    replicationState.isMaster = true
    config.replicaof = undefined
    console.log("isMaster", replicationState.isMaster)
    console.log("replicaOf", config.replicaof)
    return
  }
  if (!replicationState.isMaster) {
    config.replicaof = `${host} ${port}`
    connectToMaster()
  }
}

function handlePsync(args, socket) {
  if (args.length !== 3) throw new Error("Invalid PSYNC command")

  const [_, replicaReplId, replicaOffset] = args
  if (
    replicationState.id !== replicaReplId ||
    !isIncludedInBacklog(replicaOffset)
  ) {
    const response = serialize(
      `FULLRESYNC ${replicationState.id} ${replicationState.offset}`
    )
    const buffer = offloadRDBSave()
    updateReplicaOffset(socket, replicationState.offset)
    return { response, rdb: buffer }
  }

  if (replicaOffset !== replicationState.offset) {
    const missingBytes = replicationState.offset - replicaOffset
    const rs = readFromBacklog(replicaOffset, missingBytes)
    const commands = rs
      .split("*")
      .slice(1)
      .map((cmd) => "*" + cmd) // Remove the empty first element

    updateReplicaOffset(socket, replicationState.offset)
    return { missing: true, commands }
  }

  updateReplicaOffset(socket, replicaOffset)
  return OK
}

function offloadRDBSave() {
  const worker = new Worker("./persistence/rdb-worker.js", {
    workerData: {
      map,
      expiryMap,
    },
  })
  return new Promise((resolve, reject) => {
    worker.on("message", (data) => resolve(data))
    worker.on("error", (error) => {
      console.error("Worker error: ", error)
      reject(error)
    })
    worker.on("exit", (code) => {
      if (code !== 0) {
        console.error(`Worker stopped with exit code ${code}`)
        reject(code)
      }
    })
  })
}

function handleSave(args) {
  if (args.length !== 1) throw new Error("Invalid SAVE command")

  offloadRDBSave()
  return OK
}

function handleMulti(args, socket) {
  if (args.length !== 1) {
    throw new Error("Invalid MULTI command")
  }
  if (socket.isMultiMode) {
    return serializeError("MULTI calls cannot be nested")
  }
  socket.isMultiMode = true
  socket.isAborted = false
  socket.transactionArray = []
  return OK
}

function handleExec(args, socket) {
  if (args.length !== 1) {
    throw new Error("Invalid EXEC command")
  }
  if (!socket.isMultiMode) {
    return serializeError("EXEC without MULTI")
  }
  socket.isMultiMode = false
  if (socket.isAborted) {
    socket.transactionArray = []
    return serializeError(
      "Transaction discarded because of previous syntax errors"
    )
  }
  const result = []
  for (const command of socket.transactionArray) {
    result.push(commandHandler(command, socket))
  }
  socket.transactionArray = []
  return serialize(result)
}

function handleDiscard(args, socket) {
  if (args.length !== 1) {
    throw new Error("Invalid DISCARD command")
  }
  if (!socket.isMultiMode) {
    return serializeError("DISCARD without MULTI")
  }
  socket.isMultiMode = false
  socket.transactionArray = undefined

  return OK
}

module.exports = commandHandler