<commit_before>require 'stripmem/web'
require 'stripmem/websocket'

if RUBY_VERSION =~ /^1.8/
  module Enumerable
    def each_with_object(obj)
      each do |x|
        yield x, obj
      end
      obj
    end
  end
end

module StripMem
  class App
    def self.run!(argv)
      new(argv).run!
    end

    def initialize(argv)
      @command = argv
      @start_time = Time.now
    end

    attr_reader :command, :start_time

    def run!
      EventMachine.run do
        spawn!
        channel = EM::Channel.new
        EM::PeriodicTimer.new(1.0) { find_children }
        EM::PeriodicTimer.new(0.2) { ps(channel) }
        WebSocket.new(channel).run!
        Thread.new do
          Web.new(channel).run! # This doesn't return until sinatra exits. (Sinatra handles SIGINT.)
          kill!
          EM.stop
          exit
        end
      end
      kill!
    end

    def kill!
      Process.kill('QUIT', @child.to_i)
    rescue => e
      puts "kill #{@child.inspect}: #{e}"
    end

    def find_children
      children = `ps a -o ppid=,pid=,command=`.lines.each_with_object(Hash.new { |h,k| h[k] = [] }) { |line, h| (line =~ /(\d+) +(\d+) +(.*)/) && h[$1].push(:pid => $2, :name => $3) }
      parents = processes.keys
      while parent = parents.shift
        children[parent.to_s].each do |child|
          if child[:name] !~ /^ps /
            pid = child[:pid].to_i
            parents << pid
            processes[pid] ||= child[:name]
          end
        end
      end
    end

    def ps(channel)
      ps = `ps -o pid=,rss= -p #{processes.keys.join(',')}`.lines.each_with_object({}) { |line, h| pid, rss = line.split(/\s+/) ; h[pid.to_i] = rss.to_i }
      offset = Time.now - start_time
      channel.push :offset => offset, :samples => processes.map { |pid, name| { :name => "[#{pid}] #{name}", :rss => ps[pid] } }
    end

    def processes
      @processes ||= { $$ => '(stripmem)' }
    end

    def spawn!
      @child = fork do
        exec(*command)
      end
      processes[@child] = command.join(' ')
    end
  end
end<commit_msg>Stop sending data when the child exits.<commit_after>require 'stripmem/web'
require 'stripmem/websocket'

if RUBY_VERSION =~ /^1.8/
  module Enumerable
    def each_with_object(obj)
      each do |x|
        yield x, obj
      end
      obj
    end
  end
end

module StripMem
  class App
    def self.run!(argv)
      new(argv).run!
    end

    def initialize(argv)
      @command = argv
      @start_time = Time.now
    end

    attr_reader :command, :start_time

    def run!
      @timers = []
      EventMachine.run do
        spawn!
        channel = EM::Channel.new
        @timers = [
          EM::PeriodicTimer.new(1.0) { find_children },
          EM::PeriodicTimer.new(0.2) { ps(channel) },
          EM::PeriodicTimer.new(1.0) { wait_child },
        ]
        WebSocket.new(channel).run!
        Thread.new do
          Web.new(channel).run! # This doesn't return until sinatra exits. (Sinatra handles SIGINT.)
          kill!
          EM.stop
          exit
        end
      end
      kill!
    end

    def kill!
      Process.kill('QUIT', @child.to_i) unless @child_status
    rescue => e
      puts "kill #{@child.inspect}: #{e}"
    end

    def find_children
      children = `ps a -o ppid=,pid=,command=`.lines.each_with_object(Hash.new { |h,k| h[k] = [] }) { |line, h| (line =~ /(\d+) +(\d+) +(.*)/) && h[$1].push(:pid => $2, :name => $3) }
      parents = processes.keys
      while parent = parents.shift
        children[parent.to_s].each do |child|
          if child[:name] !~ /^ps /
            pid = child[:pid].to_i
            parents << pid
            processes[pid] ||= child[:name]
          end
        end
      end
    end

    def ps(channel)
      ps = `ps -o pid=,rss= -p #{processes.keys.join(',')}`.lines.each_with_object({}) { |line, h| pid, rss = line.split(/\s+/) ; h[pid.to_i] = rss.to_i }
      offset = Time.now - start_time
      channel.push :offset => offset, :samples => processes.map { |pid, name| { :name => "[#{pid}] #{name}", :rss => ps[pid] } }
    end

    def processes
      @processes ||= { $$ => '(stripmem)' }
    end

    def spawn!
      @child = fork do
        exec(*command)
      end
      processes[@child] = command.join(' ')
    end

    def wait_child
      if @child_status ||= Process.waitpid2(-1, Process::WNOHANG)
        @timers.each do |timer|
          timer.cancel
        end
      end
    end
  end
end
{'fraction_non_alphanumeric': 0.15078651685393257, 'fraction_numerical': 0.004943820224719101, 'mean_word_length': 1.9437830687830688, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1734360', 'n_tokens_mistral': 1537, 'n_tokens_neox': 1516, 'n_words': 362}