<commit_before>{-# LANGUAGE BangPatterns #-}

module Pygmalion.Database.Manager
( runDatabaseManager
, ensureDB
, DBRequest (..)
, DBChan
) where

import Control.Applicative
import Control.Concurrent
import Control.Monad
import Data.Time.Clock

import Control.Concurrent.Chan.Len
import Pygmalion.Core
import Pygmalion.Database.IO
import Pygmalion.Log

data DBRequest = DBUpdateCommandInfo CommandInfo
               | DBUpdateDefInfo DefInfo
               | DBUpdateOverride Override
               | DBUpdateRef Reference
               | DBUpdateInclusion Inclusion
               | DBResetMetadata SourceFile
               | DBGetCommandInfo SourceFile (MVar (Maybe CommandInfo))
               | DBGetSimilarCommandInfo SourceFile (MVar (Maybe CommandInfo))
               | DBGetDefinition USR (MVar (Maybe DefInfo))
               | DBGetIncluders SourceFile (MVar [CommandInfo])
               | DBGetCallers USR (MVar [Invocation])
               | DBGetCallees USR (MVar [DefInfo])
               | DBGetBases USR (MVar [DefInfo])
               | DBGetOverrides USR (MVar [DefInfo])
               | DBGetRefs USR (MVar [SourceReference])
               | DBGetReferenced SourceLocation (MVar [SourceReferenced])
               | DBShutdown
type DBChan = LenChan DBRequest

runDatabaseManager :: DBChan -> DBChan -> IO ()
runDatabaseManager chan queryChan = do
    start <- getCurrentTime
    withDB (go 0 start)
  where
    go :: Int -> UTCTime -> DBHandle -> IO ()
    go 1000 !start !h = do 
      stop <- getCurrentTime
      logInfo $ "Handled 1000 records in " ++ (show $ stop `diffUTCTime` start)
      newStart <- getCurrentTime
      go 0 newStart h
    go !n !s !h = {-# SCC "databaseThread" #-}
           do (!tookFirst, !newCount, !req) <- readLenChanPreferFirst queryChan chan
              logDebug $ if tookFirst then "Query channel now has " ++ (show newCount) ++ " queries waiting"
                                      else "Database channel now has " ++ (show newCount) ++ " requests waiting"
              case req of
                DBUpdateCommandInfo !ci       -> doUpdateCommandInfo h ci >> go (n+1) s h
                DBUpdateDefInfo !di           -> doUpdateDefInfo h di >> go (n+1) s h
                DBUpdateOverride !ov          -> doUpdateOverride h ov >> go (n+1) s h
                DBUpdateRef !rf               -> doUpdateRef h rf >> go (n+1) s h
                DBUpdateInclusion !ic         -> doUpdateInclusion h ic >> go (n+1) s h
                DBResetMetadata !sf           -> doResetMetadata h sf >> go (n+1) s h
                DBGetCommandInfo !f !v        -> doGetCommandInfo h f v >> go (n+1) s h
                DBGetSimilarCommandInfo !f !v -> doGetSimilarCommandInfo h f v >> go (n+1) s h
                DBGetDefinition !u !v         -> doGetDefinition h u v >> go (n+1) s h
                DBGetIncluders !sf !v         -> doGetIncluders h sf v >> go (n+1) s h
                DBGetCallers !usr !v          -> doGetCallers h usr v >> go (n+1) s h
                DBGetCallees !usr !v          -> doGetCallees h usr v >> go (n+1) s h
                DBGetBases !usr !v            -> doGetBases h usr v >> go (n+1) s h
                DBGetOverrides !usr !v        -> doGetOverrides h usr v >> go (n+1) s h
                DBGetRefs !usr !v             -> doGetRefs h usr v >> go (n+1) s h
                DBGetReferenced !sl !v        -> doGetReferenced h sl v >> go (n+1) s h
                DBShutdown                    -> logInfo "Shutting down DB thread"

doUpdateCommandInfo :: DBHandle -> CommandInfo -> IO ()
doUpdateCommandInfo h ci = withTransaction h $ do
  logDebug $ "Updating database with command: " ++ (show . ciSourceFile $ ci)
  updateSourceFile h ci

doUpdateDefInfo :: DBHandle -> DefInfo -> IO ()
doUpdateDefInfo h di = withTransaction h $ do<commit_msg>Ensure the query channel does not overwhelm the regular database channel<commit_after>{-# LANGUAGE BangPatterns #-}

module Pygmalion.Database.Manager
( runDatabaseManager
, ensureDB
, DBRequest (..)
, DBChan
) where

import Control.Applicative
import Control.Concurrent
import Control.Monad
import Data.Time.Clock

import Control.Concurrent.Chan.Len
import Pygmalion.Core
import Pygmalion.Database.IO
import Pygmalion.Log

data DBRequest = DBUpdateCommandInfo CommandInfo
               | DBUpdateDefInfo DefInfo
               | DBUpdateOverride Override
               | DBUpdateRef Reference
               | DBUpdateInclusion Inclusion
               | DBResetMetadata SourceFile
               | DBGetCommandInfo SourceFile (MVar (Maybe CommandInfo))
               | DBGetSimilarCommandInfo SourceFile (MVar (Maybe CommandInfo))
               | DBGetDefinition USR (MVar (Maybe DefInfo))
               | DBGetIncluders SourceFile (MVar [CommandInfo])
               | DBGetCallers USR (MVar [Invocation])
               | DBGetCallees USR (MVar [DefInfo])
               | DBGetBases USR (MVar [DefInfo])
               | DBGetOverrides USR (MVar [DefInfo])
               | DBGetRefs USR (MVar [SourceReference])
               | DBGetReferenced SourceLocation (MVar [SourceReferenced])
               | DBShutdown
type DBChan = LenChan DBRequest

runDatabaseManager :: DBChan -> DBChan -> IO ()
runDatabaseManager chan queryChan = do
    start <- getCurrentTime
    withDB (go 0 start)
  where
    go :: Int -> UTCTime -> DBHandle -> IO ()
    go 1000 !start !h = do 
      stop <- getCurrentTime
      logInfo $ "Handled 1000 records in " ++ (show $ stop `diffUTCTime` start)
      newStart <- getCurrentTime
      go 0 newStart h
    go !n !s !h = {-# SCC "databaseThread" #-}
           do (!tookFirst, !newCount, !req) <- readEitherChan n queryChan chan
              logDebug $ if tookFirst then "Query channel now has " ++ (show newCount) ++ " queries waiting"
                                      else "Database channel now has " ++ (show newCount) ++ " requests waiting"
              case req of
                DBUpdateCommandInfo !ci       -> doUpdateCommandInfo h ci >> go (n+1) s h
                DBUpdateDefInfo !di           -> doUpdateDefInfo h di >> go (n+1) s h
                DBUpdateOverride !ov          -> doUpdateOverride h ov >> go (n+1) s h
                DBUpdateRef !rf               -> doUpdateRef h rf >> go (n+1) s h
                DBUpdateInclusion !ic         -> doUpdateInclusion h ic >> go (n+1) s h
                DBResetMetadata !sf           -> doResetMetadata h sf >> go (n+1) s h
                DBGetCommandInfo !f !v        -> doGetCommandInfo h f v >> go (n+1) s h
                DBGetSimilarCommandInfo !f !v -> doGetSimilarCommandInfo h f v >> go (n+1) s h
                DBGetDefinition !u !v         -> doGetDefinition h u v >> go (n+1) s h
                DBGetIncluders !sf !v         -> doGetIncluders h sf v >> go (n+1) s h
                DBGetCallers !usr !v          -> doGetCallers h usr v >> go (n+1) s h
                DBGetCallees !usr !v          -> doGetCallees h usr v >> go (n+1) s h
                DBGetBases !usr !v            -> doGetBases h usr v >> go (n+1) s h
                DBGetOverrides !usr !v        -> doGetOverrides h usr v >> go (n+1) s h
                DBGetRefs !usr !v             -> doGetRefs h usr v >> go (n+1) s h
                DBGetReferenced !sl !v        -> doGetReferenced h sl v >> go (n+1) s h
                DBShutdown                    -> logInfo "Shutting down DB thread"

readEitherChan :: Int -> DBChan -> DBChan -> IO (Bool, Int, DBRequest)
readEitherChan n queryChan chan
  | n `rem` 10 == 0 = readLenChanPreferFirst queryChan chan
  | otherwise       = readLenChanPreferFirst chan queryChan

doUpdateCommandInfo :: DBHandle -> CommandInfo -> IO ()
doUpdateCommandInfo h ci = withTransaction h $ do
  logDebug $ "Updating database with command: " ++ (show . ciSourceFile $ ci)
  updateSourceFile h ci

doUpdateDefInfo :: DBHandle -> DefInfo -> IO ()
doUpdateDefInfo h di = withTransaction h $ do
{'fraction_non_alphanumeric': 0.0889558989822842, 'fraction_numerical': 0.006910415881392135, 'mean_word_length': 1.9394387001477105, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 122, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4208004', 'n_tokens_mistral': 2409, 'n_tokens_neox': 2277, 'n_words': 769}