<commit_before>from __future__ import absolute_import, division, print_function

# ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

from unittest import TestCase, main

from skbio.core.distance import DistanceMatrix
from skbio.core.tree import TreeNode
from skbio.core.tree.nj import (nj, _compute_q, _compute_collapsed_dm,
    _lowest_index, _otu_to_new_node, _pair_members_to_new_node)

class NjTests(TestCase):

    def setUp(self):
        data1 = [[0,  5,  9,  9,  8],
                 [5,  0, 10, 10,  9],
                 [9, 10,  0,  8,  7],
                 [9, 10,  8,  0,  3],
                 [8,  9,  7,  3,  0]]
        ids1 = list('abcde')
        self.dm1 = DistanceMatrix(data1, ids1)
        # this newick string was confirmed against http://www.trex.uqam.ca/
        # which generated the following (isomorphic) newick string:
        # (d:2.0000,e:1.0000,(c:4.0000,(a:2.0000,b:3.0000):3.0000):2.0000);
        self.expected1_str = ("(d:2.000000, (c:4.000000, (b:3.000000,"
        " a:2.000000):3.000000):2.000000, e:1.000000);")
        self.expected1_TreeNode = TreeNode.from_newick(self.expected1_str)

        # this example was pulled from the Phylip manual
        # http://evolution.genetics.washington.edu/phylip/doc/neighbor.html
        data2 = [[0.0000, 1.6866, 1.7198, 1.6606, 1.5243, 1.6043, 1.5905],
                 [1.6866, 0.0000, 1.5232, 1.4841, 1.4465, 1.4389, 1.4629],
                 [1.7198, 1.5232, 0.0000, 0.7115, 0.5958, 0.6179, 0.5583],
                 [1.6606, 1.4841, 0.7115, 0.0000, 0.4631, 0.5061, 0.4710],
                 [1.5243, 1.4465, 0.5958, 0.4631, 0.0000, 0.3484, 0.3083],
                 [1.6043, 1.4389, 0.6179, 0.5061, 0.3484, 0.0000, 0.2692],
                 [1.5905, 1.4629, 0.5583, 0.4710, 0.3083, 0.2692, 0.0000]]
        ids2 = ["Bovine", "Mouse", "Gibbon", "Orang", "Gorilla", "Chimp",
                "Human"]
        self.dm2 = DistanceMatrix(data2, ids2)
        self.expected2_str = ("(Mouse:0.76891, (Gibbon:0.35793, (Orang:0.28469,"
            " (Gorilla:0.15393, (Chimp:0.15167, Human:0.11753):0.03982):"
            "0.02696):0.04648):0.42027, Bovine:0.91769);")
        self.expected2_TreeNode = TreeNode.from_newick(self.expected2_str)

    def test_nj_dm1(self):
        self.assertEqual(nj(self.dm1, result_constructor=str),
                         self.expected1_str)
        # what is the correct way to compare TreeNode objects for equality?
        actual_TreeNode = nj(self.dm1)
        self.assertEqual(\
            actual_TreeNode.compare_tip_distances(self.expected1_TreeNode), 0.0)

    def test_nj_dm2(self):
        actual_TreeNode = nj(self.dm2)
        self.assertAlmostEqual(\
            actual_TreeNode.compare_tip_distances(self.expected2_TreeNode), 0.0)

    def test_nj_trivial(self):
        data = [[0, 3, 2],
                [3, 0, 3],
                [2, 3, 0]]
        dm = DistanceMatrix(data, list('abc'))
        expected_str = "(b:2.000000, a:1.000000, c:1.000000);"
        self.assertEqual(nj(dm, result_constructor=str), expected_str)

    def test_nj_error(self):
        data = [[0, 3],
                [3, 0]]
        dm = DistanceMatrix(data, list('ab'))
        self.assertRaises(ValueError, nj, dm)

    def test_compute_q(self):
        expected_data = [[  0, -50, -38, -34, -34],
                         [-50,   0, -38, -34, -34],
                         [-38, -38,   0, -40, -40],
                         [-34, -34, -40,   0, -48],
                         [-34, -34, -40, -48,   0]]
        expected_ids = list('abcde')
        expected = DistanceMatrix(expected_data, expected_ids)
        self.assertEqual(_compute_q(self.dm1), expected)

        data = [[0, 3, 2],
                [3, 0, 3],
                [2, 3, 0]]
        dm = DistanceMatrix(data, list('abc'))
        # computed this manually
        expected_data = [[ 0, -8, -8],
                         [-8,  0, -8],
                         [-8, -8,  0]]
        expected = DistanceMatrix(expected_data, list('abc'))
        self.assertEqual(_compute_q(dm), expected)

    def test_compute_collapsed_dm(self):
        expected_data = [[0,  7,  7,  6],
                         [7,  0,  8,  7],
                         [7,  8,  0,  3],
                         [6,  7,  3,  0]]
        expected_ids = ['x', 'c', 'd', 'e']
        expected1 = DistanceMatrix(expected_data, expected_ids)
        self.assertEqual(_compute_collapsed_dm(self.dm1, 'a', 'b', True, 'x'),
                         expected1)

        # computed manually
        expected_data = [[0, 4, 3],
                         [4, 0, 3],
                         [3, 3, 0]]
        expected_ids = ['yy', 'd', 'e']
        expected2 = DistanceMatrix(expected_data, expected_ids)
        self.assertEqual(_compute_collapsed_dm(expected1, 'x', 'c', True, 'yy'),
                         expected2)

    def test_lowest_index(self):
        self.assertEqual(_lowest_index(self.dm1), (4, 3))
        self.assertEqual(_lowest_index(_compute_q(self.dm1)), (1, 0))

    def test_otu_to_new_node(self):
        self.assertEqual(_otu_to_new_node(self.dm1, 'a', 'b', 'c', True), 7)

    def test_pair_members_to_new_node(self):
        self.assertEqual(_pair_members_to_new_node(self.dm1, 'a', 'b', True),
                         (2, 3))

if __name__ == "__main__":
    main()
<commit_msg>ENH: expanded tests<commit_after>from __future__ import absolute_import, division, print_function

# ----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
# ----------------------------------------------------------------------------

from unittest import TestCase, main

from skbio.core.distance import DistanceMatrix
from skbio.core.tree import TreeNode
from skbio.core.tree.nj import (nj, _compute_q, _compute_collapsed_dm,
    _lowest_index, _otu_to_new_node, _pair_members_to_new_node)

class NjTests(TestCase):

    def setUp(self):
        data1 = [[0,  5,  9,  9,  8],
                 [5,  0, 10, 10,  9],
                 [9, 10,  0,  8,  7],
                 [9, 10,  8,  0,  3],
                 [8,  9,  7,  3,  0]]
        ids1 = list('abcde')
        self.dm1 = DistanceMatrix(data1, ids1)
        # this newick string was confirmed against http://www.trex.uqam.ca/
        # which generated the following (isomorphic) newick string:
        # (d:2.0000,e:1.0000,(c:4.0000,(a:2.0000,b:3.0000):3.0000):2.0000);
        self.expected1_str = ("(d:2.000000, (c:4.000000, (b:3.000000,"
        " a:2.000000):3.000000):2.000000, e:1.000000);")
        self.expected1_TreeNode = TreeNode.from_newick(self.expected1_str)

        # this example was pulled from the Phylip manual
        # http://evolution.genetics.washington.edu/phylip/doc/neighbor.html
        data2 = [[0.0000, 1.6866, 1.7198, 1.6606, 1.5243, 1.6043, 1.5905],
                 [1.6866, 0.0000, 1.5232, 1.4841, 1.4465, 1.4389, 1.4629],
                 [1.7198, 1.5232, 0.0000, 0.7115, 0.5958, 0.6179, 0.5583],
                 [1.6606, 1.4841, 0.7115, 0.0000, 0.4631, 0.5061, 0.4710],
                 [1.5243, 1.4465, 0.5958, 0.4631, 0.0000, 0.3484, 0.3083],
                 [1.6043, 1.4389, 0.6179, 0.5061, 0.3484, 0.0000, 0.2692],
                 [1.5905, 1.4629, 0.5583, 0.4710, 0.3083, 0.2692, 0.0000]]
        ids2 = ["Bovine", "Mouse", "Gibbon", "Orang", "Gorilla", "Chimp",
                "Human"]
        self.dm2 = DistanceMatrix(data2, ids2)
        self.expected2_str = ("(Mouse:0.76891, (Gibbon:0.35793, (Orang:0.28469,"
            " (Gorilla:0.15393, (Chimp:0.15167, Human:0.11753):0.03982):"
            "0.02696):0.04648):0.42027, Bovine:0.91769);")
        self.expected2_TreeNode = TreeNode.from_newick(self.expected2_str)

    def test_nj_dm1(self):
        self.assertEqual(nj(self.dm1, result_constructor=str),
                         self.expected1_str)
        # what is the correct way to compare TreeNode objects for equality?
        actual_TreeNode = nj(self.dm1)
        self.assertEqual(\
            actual_TreeNode.compare_tip_distances(self.expected1_TreeNode), 0.0)

    def test_nj_dm2(self):
        actual_TreeNode = nj(self.dm2)
        self.assertAlmostEqual(\
            actual_TreeNode.compare_tip_distances(self.expected2_TreeNode), 0.0)

    def test_nj_trivial(self):
        data = [[0, 3, 2],
                [3, 0, 3],
                [2, 3, 0]]
        dm = DistanceMatrix(data, list('abc'))
        expected_str = "(b:2.000000, a:1.000000, c:1.000000);"
        self.assertEqual(nj(dm, result_constructor=str), expected_str)

    def test_nj_error(self):
        data = [[0, 3],
                [3, 0]]
        dm = DistanceMatrix(data, list('ab'))
        self.assertRaises(ValueError, nj, dm)

    def test_compute_q(self):
        expected_data = [[  0, -50, -38, -34, -34],
                         [-50,   0, -38, -34, -34],
                         [-38, -38,   0, -40, -40],
                         [-34, -34, -40,   0, -48],
                         [-34, -34, -40, -48,   0]]
        expected_ids = list('abcde')
        expected = DistanceMatrix(expected_data, expected_ids)
        self.assertEqual(_compute_q(self.dm1), expected)

        data = [[0, 3, 2],
                [3, 0, 3],
                [2, 3, 0]]
        dm = DistanceMatrix(data, list('abc'))
        # computed this manually
        expected_data = [[ 0, -8, -8],
                         [-8,  0, -8],
                         [-8, -8,  0]]
        expected = DistanceMatrix(expected_data, list('abc'))
        self.assertEqual(_compute_q(dm), expected)

    def test_compute_collapsed_dm(self):
        expected_data = [[0,  7,  7,  6],
                         [7,  0,  8,  7],
                         [7,  8,  0,  3],
                         [6,  7,  3,  0]]
        expected_ids = ['x', 'c', 'd', 'e']
        expected1 = DistanceMatrix(expected_data, expected_ids)
        self.assertEqual(_compute_collapsed_dm(self.dm1, 'a', 'b', True, 'x'),
                         expected1)

        # computed manually
        expected_data = [[0, 4, 3],
                         [4, 0, 3],
                         [3, 3, 0]]
        expected_ids = ['yy', 'd', 'e']
        expected2 = DistanceMatrix(expected_data, expected_ids)
        self.assertEqual(_compute_collapsed_dm(expected1, 'x', 'c', True, 'yy'),
                         expected2)

    def test_lowest_index(self):
        self.assertEqual(_lowest_index(self.dm1), (4, 3))
        self.assertEqual(_lowest_index(_compute_q(self.dm1)), (1, 0))

    def test_otu_to_new_node(self):
        self.assertEqual(_otu_to_new_node(self.dm1, 'a', 'b', 'c', True), 7)
        self.assertEqual(_otu_to_new_node(self.dm1, 'a', 'b', 'd', True), 7)
        self.assertEqual(_otu_to_new_node(self.dm1, 'a', 'b', 'e', True), 6)

    def test_pair_members_to_new_node(self):
        self.assertEqual(_pair_members_to_new_node(self.dm1, 'a', 'b', True),
                         (2, 3))
        self.assertEqual(_pair_members_to_new_node(self.dm1, 'a', 'c', True),
                         (4, 5))
        self.assertEqual(_pair_members_to_new_node(self.dm1, 'd', 'e', True),
                         (2, 1))

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before># Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
import errno
import json
import logging
import requests
import threading
import time
import websocket

import parlai.mturk.core.shared_utils as shared_utils

MAX_QUICK_REPLIES = 10
MAX_TEXT_CHARS = 640
MAX_QUICK_REPLY_TITLE_CHARS = 20
MAX_POSTBACK_CHARS = 1000
SOCKET_TIMEOUT = 6

# Arbitrary attachments can be created as long as they adhere to the docs
# developers.facebook.com/docs/messenger-platform/send-messages/templates

# Message builders
def create_attachment(attachment_type, url):
    """Create a simple url-based attachment"""
    # TODO support data-based attachments?
    assert attachment_type in ['image', 'video', 'file', 'audio'], \
        'unsupported attachment type'
    return {'type': attachment_type, 'payload': {'url': url}}

def create_reply_option(title, payload=''):
    """Create a quick reply option. Takes in display title and optionally extra
    custom data.
    """
    assert len(title) <= MAX_QUICK_REPLY_TITLE_CHARS, (
        'Quick reply title length {} greater than the max of {}'.format(
            len(title), MAX_QUICK_REPLY_TITLE_CHARS
        )
    )
    assert len(payload) <= MAX_POSTBACK_CHARS, (
        'Payload length {} greater than the max of {}'.format(
            len(payload), MAX_POSTBACK_CHARS
        )
    )
    return {'content_type': 'text', 'title': title, 'payload': payload}

def create_text_message(text, quick_replies=None):
    """Return a list of text messages from the given text. If the message is
    too long it is split into multiple messages.
    quick_replies should be a list of options made with create_reply_option.
    """
    def _message(text_content, replies):
        payload = {'text': text_content[:MAX_TEXT_CHARS]}
        if replies:
            payload['quick_replies'] = replies
        return payload

    tokens = [s[:MAX_TEXT_CHARS] for s in text.split(' ')]
    splits = []
    cutoff = 0
    curr_length = 0
    if quick_replies:
        assert len(quick_replies) <= MAX_QUICK_REPLIES, (
            'Number of quick replies {} greater than the max of {}'.format(
                len(quick_replies), MAX_QUICK_REPLIES
            )
        )
    for i in range(len(tokens)):
        if (curr_length + len(tokens[i]) > MAX_TEXT_CHARS):
            splits.append(_message(' '.join(tokens[cutoff:i]), None))
            cutoff = i + 1
            curr_length = 0
        curr_length += len(tokens[i]) + 1
    if cutoff < len(tokens):
        splits.append(_message(' '.join(tokens[cutoff:]), quick_replies))
    return splits

def create_attachment_message(attachment_item, quick_replies=None):
    """Create a message list made with only an attachment.
    quick_replies should be a list of options made with create_reply_option.
    """
    payload = {'attachment': attachment_item}
    if quick_replies:
        assert len(quick_replies) <= MAX_QUICK_REPLIES, (
            'Number of quick replies {} greater than the max of {}'.format(
                len(quick_replies), MAX_QUICK_REPLIES
            )
        )
        payload['quick_replies'] = quick_replies
    return [payload]

def create_list_element(element):
    assert 'title' in element, 'List elems must have a title'
    ret_elem = {
        'title': element['title'],
        'subtitle': '',
        'default_action': {
            'type': 'postback',
            'title': element['title'],
            'payload': element['title'],
        }
    }
    if 'subtitle' in element:
        ret_elem['subtitle'] = element['subtitle']
    return ret_elem

def create_compact_list_message(raw_elems):
    elements = [create_list_element(elem) for elem in raw_elems]
    return {
        'type': 'template',
        'payload': {
            'template_type': 'list',
            'top_element_style': 'COMPACT',
            'elements': elements,
        }
    }

# Socket handler
class MessageSocket():
    """MessageSocket is a wrapper around websocket to simplify message sends
    and recieves into parlai from FB messenger.
    """

    def __init__(self, server_url, port, secret_token, message_callback):
        """
        server_url:           url at which the server is to be run
        port:                 port for the socket to operate on
        message_callback:     function to be called on incoming message objects
                              format: message_callback(self, data)
        """
        self.server_url = server_url
        self.port = port
        self.message_callback = message_callback

        self.ws = None
        self.last_pong = None
        self.alive = False
        self.auth_args = {'access_token': secret_token}

        # initialize the state
        self.listen_thread = None

        # setup the socket
        self.keep_running = True
        self._setup_socket()

    def _safe_send(self, data, force=False):
        if not self.alive and not force:
            # Try to wait a second to send a packet
            timeout = 1
            while timeout > 0 and not self.alive:
                time.sleep(0.1)
                timeout -= 0.1
            if not self.alive:
                # don't try to send a packet if we're still dead
                return False
        try:
            self.ws.send(data)
        except websocket.WebSocketConnectionClosedException:
            # The channel died mid-send, wait for it to come back up
            return False
        return True

    def _ensure_closed(self):
        try:
            self.ws.close()
        except websocket.WebSocketConnectionClosedException:
            pass

    def _send_world_alive(self):
        """Registers world with the passthrough server"""
        self._safe_send(json.dumps({
            'type': 'world_alive',
            'content': {'id': 'WORLD_ALIVE', 'sender_id': 'world'},
        }), force=True)

    def send_fb_payload(self, receiver_id, payload):
        """Sends a payload to messenger, processes it if we can"""
        api_address = 'https://graph.facebook.com/v2.6/me/messages'
        if payload['type'] == 'list':
            data = create_compact_list_message(payload['data'])
        else:
            data = payload['data']
        message = {
            "messaging_type": 'RESPONSE',
            "recipient": {
                "id": receiver_id
            },
            "message": {
                "attachment": data,
            }
        }
        response = requests.post(
            api_address,
            params=self.auth_args,
            json=message,
        )
        result = response.json()
        shared_utils.print_and_log(
            logging.INFO,
            '"Facebook response from message send: {}"'.format(result)
        )
        return result

    def send_fb_message(self, receiver_id, message, is_response,
                        quick_replies=None):
        """Sends a message directly to messenger"""
        api_address = 'https://graph.facebook.com/v2.6/me/messages'
        if quick_replies is not None:
            quick_replies = [create_reply_option(x, x) for x in quick_replies]
        ms = create_text_message(message, quick_replies)
        results = []
        for m in ms:
            if m['text'] == '':
                continue  # Skip blank messages
            payload = {
                "messaging_type": 'RESPONSE' if is_response else 'UPDATE',
                "recipient": {
                    "id": receiver_id
                },
                "message": m
            }
            response = requests.post(
                api_address,
                params=self.auth_args,
                json=payload
            )
            result = response.json()
            shared_utils.print_and_log(
                logging.INFO,
                '"Facebook response from message send: {}"'.format(result)
            )
            results.append(result)
        return results

    def _setup_socket(self):
        """Create socket handlers and registers the socket"""
        def on_socket_open(*args):
            shared_utils.print_and_log(
                logging.DEBUG,
                'Socket open: {}'.format(args)
            )
            self._send_world_alive()

        def on_error(ws, error):
            try:
                if error.errno == errno.ECONNREFUSED:
                    self._ensure_closed()
                    self.use_socket = False
                    raise Exception("Socket refused connection, cancelling")
                else:
                    shared_utils.print_and_log(
                        logging.WARN,
                        'Socket logged error: {}'.format(repr(error)),
                    )
            except BaseException:
                if type(error) is websocket.WebSocketConnectionClosedException:
                    return  # Connection closed is noop
                shared_utils.print_and_log(
                    logging.WARN,
                    'Socket logged error: {} Restarting'.format(repr(error)),
                )
                self._ensure_closed()

        def on_disconnect(*args):
            """Disconnect event is a no-op for us, as the server reconnects
            automatically on a retry"""
            shared_utils.print_and_log(
                logging.INFO,
                'World server disconnected: {}'.format(args)
            )
            self.alive = False
            self._ensure_closed()

        def on_message(*args):
            """Incoming message handler for messages from the FB user"""
            packet_dict = json.loads(args[1])
            if packet_dict['type'] == 'conn_success':
                self.alive = True
                return  # No action for successful connection
            if packet_dict['type'] == 'pong':
                self.last_pong = time.time()
                return  # No further action for pongs
            message_data = packet_dict['content']
            shared_utils.print_and_log(
                logging.DEBUG,
                'Message data recieved: {}'.format(message_data)
            )
            for message_packet in message_data['entry']:
                self.message_callback(message_packet['messaging'][0])

        def run_socket(*args):
            url_base_name = self.server_url.split('https://')[1]
            while self.keep_running:
                try:
                    sock_addr = "ws://{}/".format(
                        url_base_name)
                    self.ws = websocket.WebSocketApp(
                        sock_addr,
                        on_message=on_message,
                        on_error=on_error,
                        on_close=on_disconnect,
                    )
                    self.ws.on_open = on_socket_open
                    self.ws.run_forever(ping_interval=1, ping_timeout=0.9)
                except Exception as e:
                    shared_utils.print_and_log(
                        logging.WARN,
                        'Socket error {}, attempting restart'.format(repr(e))
                    )
                time.sleep(0.2)

        # Start listening thread
        self.listen_thread = threading.Thread(
            target=run_socket,
            name='Main-Socket-Thread'
        )
        self.listen_thread.daemon = True
        self.listen_thread.start()
        time.sleep(1.2)
        while not self.alive:
            try:
                self._send_world_alive()
            except Exception:
                pass
            time.sleep(0.8)
<commit_msg>expose way to send media through ParlAI messenger<commit_after># Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
import errno
import json
import logging
import requests
import threading
import time
import websocket

import parlai.mturk.core.shared_utils as shared_utils

MAX_QUICK_REPLIES = 10
MAX_TEXT_CHARS = 640
MAX_QUICK_REPLY_TITLE_CHARS = 20
MAX_POSTBACK_CHARS = 1000
SOCKET_TIMEOUT = 6

# Arbitrary attachments can be created as long as they adhere to the docs
# developers.facebook.com/docs/messenger-platform/send-messages/templates

# Message builders
def create_attachment(attachment_type, url):
    """Create a simple url-based attachment"""
    # TODO support data-based attachments?
    assert attachment_type in ['image', 'video', 'file', 'audio'], \
        'unsupported attachment type'
    return {'type': attachment_type, 'payload': {'url': url}}

def create_reply_option(title, payload=''):
    """Create a quick reply option. Takes in display title and optionally extra
    custom data.
    """
    assert len(title) <= MAX_QUICK_REPLY_TITLE_CHARS, (
        'Quick reply title length {} greater than the max of {}'.format(
            len(title), MAX_QUICK_REPLY_TITLE_CHARS
        )
    )
    assert len(payload) <= MAX_POSTBACK_CHARS, (
        'Payload length {} greater than the max of {}'.format(
            len(payload), MAX_POSTBACK_CHARS
        )
    )
    return {'content_type': 'text', 'title': title, 'payload': payload}

def create_text_message(text, quick_replies=None):
    """Return a list of text messages from the given text. If the message is
    too long it is split into multiple messages.
    quick_replies should be a list of options made with create_reply_option.
    """
    def _message(text_content, replies):
        payload = {'text': text_content[:MAX_TEXT_CHARS]}
        if replies:
            payload['quick_replies'] = replies
        return payload

    tokens = [s[:MAX_TEXT_CHARS] for s in text.split(' ')]
    splits = []
    cutoff = 0
    curr_length = 0
    if quick_replies:
        assert len(quick_replies) <= MAX_QUICK_REPLIES, (
            'Number of quick replies {} greater than the max of {}'.format(
                len(quick_replies), MAX_QUICK_REPLIES
            )
        )
    for i in range(len(tokens)):
        if (curr_length + len(tokens[i]) > MAX_TEXT_CHARS):
            splits.append(_message(' '.join(tokens[cutoff:i]), None))
            cutoff = i + 1
            curr_length = 0
        curr_length += len(tokens[i]) + 1
    if cutoff < len(tokens):
        splits.append(_message(' '.join(tokens[cutoff:]), quick_replies))
    return splits

def create_attachment_message(attachment_item, quick_replies=None):
    """Create a message list made with only an attachment.
    quick_replies should be a list of options made with create_reply_option.
    """
    payload = {'attachment': attachment_item}
    if quick_replies:
        assert len(quick_replies) <= MAX_QUICK_REPLIES, (
            'Number of quick replies {} greater than the max of {}'.format(
                len(quick_replies), MAX_QUICK_REPLIES
            )
        )
        payload['quick_replies'] = quick_replies
    return [payload]

def create_list_element(element):
    assert 'title' in element, 'List elems must have a title'
    ret_elem = {
        'title': element['title'],
        'subtitle': '',
        'default_action': {
            'type': 'postback',
            'title': element['title'],
            'payload': element['title'],
        }
    }
    if 'subtitle' in element:
        ret_elem['subtitle'] = element['subtitle']
    return ret_elem

def create_compact_list_message(raw_elems):
    elements = [create_list_element(elem) for elem in raw_elems]
    return {
        'type': 'template',
        'payload': {
            'template_type': 'list',
            'top_element_style': 'COMPACT',
            'elements': elements,
        }
    }

# Socket handler
class MessageSocket():
    """MessageSocket is a wrapper around websocket to simplify message sends
    and recieves into parlai from FB messenger.
    """

    def __init__(self, server_url, port, secret_token, message_callback):
        """
        server_url:           url at which the server is to be run
        port:                 port for the socket to operate on
        message_callback:     function to be called on incoming message objects
                              format: message_callback(self, data)
        """
        self.server_url = server_url
        self.port = port
        self.message_callback = message_callback

        self.ws = None
        self.last_pong = None
        self.alive = False
        self.auth_args = {'access_token': secret_token}

        # initialize the state
        self.listen_thread = None

        # setup the socket
        self.keep_running = True
        self._setup_socket()

    def _safe_send(self, data, force=False):
        if not self.alive and not force:
            # Try to wait a second to send a packet
            timeout = 1
            while timeout > 0 and not self.alive:
                time.sleep(0.1)
                timeout -= 0.1
            if not self.alive:
                # don't try to send a packet if we're still dead
                return False
        try:
            self.ws.send(data)
        except websocket.WebSocketConnectionClosedException:
            # The channel died mid-send, wait for it to come back up
            return False
        return True

    def _ensure_closed(self):
        try:
            self.ws.close()
        except websocket.WebSocketConnectionClosedException:
            pass

    def _send_world_alive(self):
        """Registers world with the passthrough server"""
        self._safe_send(json.dumps({
            'type': 'world_alive',
            'content': {'id': 'WORLD_ALIVE', 'sender_id': 'world'},
        }), force=True)

    def send_fb_payload(self, receiver_id, payload):
        """Sends a payload to messenger, processes it if we can"""
        api_address = 'https://graph.facebook.com/v2.6/me/messages'
        if payload['type'] == 'list':
            data = create_compact_list_message(payload['data'])
        elif payload['type'] in ['image', 'video', 'file', 'audio']:
            data = create_attachment(payload['type'], payload['url'])
        else:
            data = payload['data']
        message = {
            "messaging_type": 'RESPONSE',
            "recipient": {
                "id": receiver_id
            },
            "message": {
                "attachment": data,
            }
        }
        response = requests.post(
            api_address,
            params=self.auth_args,
            json=message,
        )
        result = response.json()
        shared_utils.print_and_log(
            logging.INFO,
            '"Facebook response from message send: {}"'.format(result)
        )
        return result

    def send_fb_message(self, receiver_id, message, is_response,
                        quick_replies=None):
        """Sends a message directly to messenger"""
        api_address = 'https://graph.facebook.com/v2.6/me/messages'
        if quick_replies is not None:
            quick_replies = [create_reply_option(x, x) for x in quick_replies]
        ms = create_text_message(message, quick_replies)
        results = []
        for m in ms:
            if m['text'] == '':
                continue  # Skip blank messages
            payload = {
                "messaging_type": 'RESPONSE' if is_response else 'UPDATE',
                "recipient": {
                    "id": receiver_id
                },
                "message": m
            }
            response = requests.post(
                api_address,
                params=self.auth_args,
                json=payload
            )
            result = response.json()
            shared_utils.print_and_log(
                logging.INFO,
                '"Facebook response from message send: {}"'.format(result)
            )
            results.append(result)
        return results

    def _setup_socket(self):
        """Create socket handlers and registers the socket"""
        def on_socket_open(*args):
            shared_utils.print_and_log(
                logging.DEBUG,
                'Socket open: {}'.format(args)
            )
            self._send_world_alive()

        def on_error(ws, error):
            try:
                if error.errno == errno.ECONNREFUSED:
                    self._ensure_closed()
                    self.use_socket = False
                    raise Exception("Socket refused connection, cancelling")
                else:
                    shared_utils.print_and_log(
                        logging.WARN,
                        'Socket logged error: {}'.format(repr(error)),
                    )
            except BaseException:
                if type(error) is websocket.WebSocketConnectionClosedException:
                    return  # Connection closed is noop
                shared_utils.print_and_log(
                    logging.WARN,
                    'Socket logged error: {} Restarting'.format(repr(error)),
                )
                self._ensure_closed()

        def on_disconnect(*args):
            """Disconnect event is a no-op for us, as the server reconnects
            automatically on a retry"""
            shared_utils.print_and_log(
                logging.INFO,
                'World server disconnected: {}'.format(args)
            )
            self.alive = False
            self._ensure_closed()

        def on_message(*args):
            """Incoming message handler for messages from the FB user"""
            packet_dict = json.loads(args[1])
            if packet_dict['type'] == 'conn_success':
                self.alive = True
                return  # No action for successful connection
            if packet_dict['type'] == 'pong':
                self.last_pong = time.time()
                return  # No further action for pongs
            message_data = packet_dict['content']
            shared_utils.print_and_log(
                logging.DEBUG,
                'Message data recieved: {}'.format(message_data)
            )
            for message_packet in message_data['entry']:
                self.message_callback(message_packet['messaging'][0])

        def run_socket(*args):
            url_base_name = self.server_url.split('https://')[1]
            while self.keep_running:
                try:
                    sock_addr = "ws://{}/".format(
                        url_base_name)
                    self.ws = websocket.WebSocketApp(
                        sock_addr,
                        on_message=on_message,
                        on_error=on_error,
                        on_close=on_disconnect,
                    )
                    self.ws.on_open = on_socket_open
                    self.ws.run_forever(ping_interval=1, ping_timeout=0.9)
                except Exception as e:
                    shared_utils.print_and_log(
                        logging.WARN,
                        'Socket error {}, attempting restart'.format(repr(e))
                    )
                time.sleep(0.2)

        # Start listening thread
        self.listen_thread = threading.Thread(
            target=run_socket,
            name='Main-Socket-Thread'
        )
        self.listen_thread.daemon = True
        self.listen_thread.start()
        time.sleep(1.2)
        while not self.alive:
            try:
                self._send_world_alive()
            except Exception:
                pass
            time.sleep(0.8)
<|endoftext|>


<commit_before>from collections import Counter
from django.shortcuts import render
from django.http import HttpResponse, HttpResponseNotFound
from django.template import loader
from django.db.models import Count
from django.core.urlresolvers import reverse

from django_peeringdb.models.concrete import *

from ipaddress import *

from prngmgr.settings import *
from prngmgr.models import *
from prngmgr.forms import *
from prngmgr.snmp import Get, GetBGPTable

me = Network.objects.get(asn=MY_ASN)

def index(request):
    template = loader.get_template('prngmgr/index.html')
    context = {
        'me': me,
    }
    return HttpResponse(template.render(context, request))

def sessions(request, filter_field, filter_arg):
    template = loader.get_template('prngmgr/table.html')
    context = {
        'view': 'sessions',
        'url_name': 'prngmgr-sessions',
        'table': {
            'name': 'sessions',
            'title': 'Peering Sessions',
            'cols': [
                { 'title': 'IXP' },
                { 'title': 'Router' },
                { 'title': 'Peer Name' },
                { 'title': 'Peer AS' },
                { 'title': 'Address Family' },
                { 'title': 'Peer Address' },
                { 'title': 'Provisioning State' },
                { 'title': 'Admin State' },
                { 'title': 'Operational State' },
            ],
            'rows': [
            ],
        },
        'filter': {
        },
    }
    if filter_field:
        context['filter'] = {
            'field': filter_field,
            'arg': filter_arg,
        }
        sessions = PeeringSession.objects.filter(**{filter_field: filter_arg})
    else:
        sessions = PeeringSession.objects.all()
    for session in sessions:
        row = { 'fields': [] }
        row['fields'].append({
            'display': session.prngrtriface.netixlan.ixlan.ix.name,
            'link': reverse('prngmgr-ixps', kwargs = { 'ixp_id': session.prngrtriface.netixlan.ixlan.ix.id })
        })
        row['fields'].append({
            'display': session.prngrtriface.prngrtr.hostname,
            'link': reverse('prngmgr-routers', kwargs = { 'rtr_id': session.prngrtriface.prngrtr.id })
        })
        row['fields'].append({
            'display': session.peer_netixlan.net.name,
            'link': reverse('prngmgr-networks', kwargs = { 'net_id': session.peer_netixlan.net.id })
        })
        row['fields'].append({
            'display': session.peer_netixlan.asn,
        })
        row['fields'].append({
            'display': session.get_af_display,
        })
        row['fields'].append({
            'display': session.get_remote_address,
        })
        row['fields'].append({
            'display': session.get_provisioning_state_display,
            'alert': session.get_provisioning_state_alert,
        })
        row['fields'].append({
            'display': session.get_admin_state_display,
            'alert': session.get_admin_state_alert,
        })
        row['fields'].append({
            'display': session.get_operational_state_display,
            'alert': session.get_operational_state_alert,
        })
        context['table']['rows'].append(row)
    return HttpResponse(template.render(context, request))

def networks(request, net_id):
    template = loader.get_template('prngmgr/table.html')
    context = {
        'view': 'networks',
        'table': {
            'name': 'networks',
            'title': 'Peering Networks',
            'cols': [
                { 'title': 'Network Name' },
                { 'title': 'Primary ASN' },
                { 'title': 'IRR Record' },
                { 'title': 'Looking Glass' },
                { 'title': 'Peering Policy' },
                { 'title': 'Possible Sessions' },
                { 'title': 'Provisioned Sessions' },
                { 'title': 'Established Sessions' },
            ],
            'rows': [
            ],
        },
        'filter': {
        },
    }
    networks = Network.objects.all()
    for network in networks:
        sessions = PeeringSession.objects.filter(peer_netixlan__net=network)
        calculated = _render_alerts({
            'count': {
                'possible': sessions.count(),
                'provisioned': sessions.filter(provisioning_state=PeeringSession.PROV_COMPLETE).count(),
                'established': sessions.filter(operational_state=PeeringSession.OPER_ESTABLISHED).count(),
            },
        })
        row = { 'fields': [] }
        row['fields'].append({
            'display': network.name,
        })
        row['fields'].append({
            'display': network.asn,
        })
        row['fields'].append({
            'display': network.irr_as_set,
        })
        row['fields'].append({
            'display': network.looking_glass,
            'link': network.looking_glass,
        })
        row['fields'].append({
            'display': network.policy_general,
        })
        row['fields'].append({
            'display': calculated['count']['possible'],
            'alert': calculated['alert']['possible'],
        })
        row['fields'].append({
            'display': calculated['count']['provisioned'],
            'alert': calculated['alert']['provisioned'],
        })
        row['fields'].append({
            'display': calculated['count']['established'],
            'alert': calculated['alert']['established'],
        })
        context['table']['rows'].append(row)
    return HttpResponse(template.render(context, request))

def routers(request, rtr_id):
    context = {
        'view': {
            'name': 'routers',
            'url': reverse('prngmgr-routers'),
        }
    }
    if request.method == 'GET':
        if rtr_id:
            template = loader.get_template('prngmgr/form.html')
            if rtr_id == 0:
                form = PeeringRouterForm()
                return HttpResponse(form)
            else:
                try:
                    router = PeeringRouter.objects.get(id=rtr_id)
                except:
                    return HttpResponseNotFound(rtr_id)
                form = PeeringRouterForm(instance=router)
            context['form'] = form
        else:
            template = loader.get_template('prngmgr/table.html')
            table = {
                'name': 'routers',
                'title': 'Peering Routers',
                'cols': [
                    { 'title': 'Hostname' },
                    { 'title': 'Peering Interfaces' },
                    { 'title': 'Possible Sessions' },
                    { 'title': 'Provisioned Sessions' },
                    { 'title': 'Established Sessions' },
                ],
                'rows': [
                ],
                'filter': {
                },
            }
            routers = PeeringRouter.objects.all()
            for router in routers:
                interfaces = PeeringRouterIXInterface.objects.filter(prngrtr=router)
                sessions = PeeringSession.objects.filter(prngrtriface__prngrtr=router)
                calculated = _render_alerts({
                    'count': {
                        'interfaces': interfaces.count(),
                        'possible': sessions.count(),
                        'provisioned': sessions.filter(provisioning_state=PeeringSession.PROV_COMPLETE).count(),
                        'established': sessions.filter(operational_state=PeeringSession.OPER_ESTABLISHED).count(),
                    },
                })
                row = { 'fields': [] }
                row['fields'].append({
                    'display': router.hostname,
                })
                row['fields'].append({
                    'display': calculated['count']['interfaces'],
                })
                row['fields'].append({
                    'display': calculated['count']['possible'],
                    'alert': calculated['alert']['possible'],
                })
                row['fields'].append({
                    'display': calculated['count']['provisioned'],
                    'alert': calculated['alert']['provisioned'],
                })
                row['fields'].append({
                    'display': calculated['count']['established'],
                    'alert': calculated['alert']['established'],
                })
                table['rows'].append(row)
            context['table'] = table
        return HttpResponse(template.render(context, request))

def ixps(request, ixp_id):
    template = loader.get_template('prngmgr/table.html')
    context = {
        'view': 'ixps',
        'table': {
            'name': 'ixps',
            'title': 'Internet Exchange Points',
            'cols': [
                { 'title': 'IXP Name' },
                { 'title': 'Country' },
                { 'title': 'Region' },
                { 'title': 'Participants' },
                { 'title': 'Possible Sessions' },
                { 'title': 'Provisioned Sessions' },
                { 'title': 'Established Sessions' },
            ],
            'rows': [
            ],
        },
        'filter': {
        },
    }
    ixps = InternetExchange.objects.all()
    for ixp in ixps:
        participants = NetworkIXLan.objects.filter(ixlan__ix=ixp).aggregate(count=Count('asn', distinct=True))
        sessions = PeeringSession.objects.filter(prngrtriface__netixlan__ixlan__ix=ixp)
        calculated = _render_alerts({
            'count': {
                'participants': participants['count'],
                'possible': sessions.count(),
                'provisioned': sessions.filter(provisioning_state=PeeringSession.PROV_COMPLETE).count(),
                'established': sessions.filter(operational_state=PeeringSession.OPER_ESTABLISHED).count(),
            },
        })
        row = { 'fields': [] }
        row['fields'].append({
            'display': ixp.name,
        })
        row['fields'].append({
            'display': ixp.country,
        })
        row['fields'].append({
            'display': ixp.region_continent,
        })
        row['fields'].append({
            'display': calculated['count']['participants'],
        })
        row['fields'].append({
            'display': calculated['count']['possible'],
            'alert': calculated['alert']['possible'],
        })
        row['fields'].append({
            'display': calculated['count']['provisioned'],
            'alert': calculated['alert']['provisioned'],
        })
        row['fields'].append({
            'display': calculated['count']['established'],
            'alert': calculated['alert']['established'],
        })
        context['table']['rows'].append(row)
    return HttpResponse(template.render(context, request))

def _render_alerts(calculated):
    if calculated['count']['possible'] == 0:
        calculated['alert'] = {
            'possible': ALERT_NONE,
            'provisioned': ALERT_NONE,
            'established': ALERT_NONE,
        }
    else:
        if calculated['count']['provisioned'] == 0:
            calculated['alert'] = {
                'possible': ALERT_SUCCESS,
                'provisioned': ALERT_DANGER,
                'established': ALERT_DANGER,
            }
        elif calculated['count']['provisioned'] < calculated['count']['possible']:
            if calculated['count']['established'] < calculated['count']['provisioned']:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_WARNING,
                    'established': ALERT_DANGER,
                }
            else:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_WARNING,
                    'established': ALERT_WARNING,
                }
        else:
            if calculated['count']['established'] < calculated['count']['provisioned']:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_SUCCESS,
                    'established': ALERT_DANGER,
                }
            else:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_SUCCESS,
                    'established': ALERT_SUCCESS,
                }
    return calculated

<commit_msg>added forms template<commit_after>from collections import Counter
from django.shortcuts import render
from django.http import HttpResponse, HttpResponseNotFound
from django.template import loader
from django.db.models import Count
from django.core.urlresolvers import reverse

from django_peeringdb.models.concrete import *

from ipaddress import *

from prngmgr.settings import *
from prngmgr.models import *
from prngmgr.forms import *
from prngmgr.snmp import Get, GetBGPTable

me = Network.objects.get(asn=MY_ASN)

def index(request):
    template = loader.get_template('prngmgr/index.html')
    context = {
        'me': me,
    }
    return HttpResponse(template.render(context, request))

def sessions(request, filter_field, filter_arg):
    template = loader.get_template('prngmgr/table.html')
    context = {
        'view': 'sessions',
        'url_name': 'prngmgr-sessions',
        'table': {
            'name': 'sessions',
            'title': 'Peering Sessions',
            'cols': [
                { 'title': 'IXP' },
                { 'title': 'Router' },
                { 'title': 'Peer Name' },
                { 'title': 'Peer AS' },
                { 'title': 'Address Family' },
                { 'title': 'Peer Address' },
                { 'title': 'Provisioning State' },
                { 'title': 'Admin State' },
                { 'title': 'Operational State' },
            ],
            'rows': [
            ],
        },
        'filter': {
        },
    }
    if filter_field:
        context['filter'] = {
            'field': filter_field,
            'arg': filter_arg,
        }
        sessions = PeeringSession.objects.filter(**{filter_field: filter_arg})
    else:
        sessions = PeeringSession.objects.all()
    for session in sessions:
        row = { 'fields': [] }
        row['fields'].append({
            'display': session.prngrtriface.netixlan.ixlan.ix.name,
            'link': reverse('prngmgr-ixps', kwargs = { 'ixp_id': session.prngrtriface.netixlan.ixlan.ix.id })
        })
        row['fields'].append({
            'display': session.prngrtriface.prngrtr.hostname,
            'link': reverse('prngmgr-routers', kwargs = { 'rtr_id': session.prngrtriface.prngrtr.id })
        })
        row['fields'].append({
            'display': session.peer_netixlan.net.name,
            'link': reverse('prngmgr-networks', kwargs = { 'net_id': session.peer_netixlan.net.id })
        })
        row['fields'].append({
            'display': session.peer_netixlan.asn,
        })
        row['fields'].append({
            'display': session.get_af_display,
        })
        row['fields'].append({
            'display': session.get_remote_address,
        })
        row['fields'].append({
            'display': session.get_provisioning_state_display,
            'alert': session.get_provisioning_state_alert,
        })
        row['fields'].append({
            'display': session.get_admin_state_display,
            'alert': session.get_admin_state_alert,
        })
        row['fields'].append({
            'display': session.get_operational_state_display,
            'alert': session.get_operational_state_alert,
        })
        context['table']['rows'].append(row)
    return HttpResponse(template.render(context, request))

def networks(request, net_id):
    template = loader.get_template('prngmgr/table.html')
    context = {
        'view': 'networks',
        'table': {
            'name': 'networks',
            'title': 'Peering Networks',
            'cols': [
                { 'title': 'Network Name' },
                { 'title': 'Primary ASN' },
                { 'title': 'IRR Record' },
                { 'title': 'Looking Glass' },
                { 'title': 'Peering Policy' },
                { 'title': 'Possible Sessions' },
                { 'title': 'Provisioned Sessions' },
                { 'title': 'Established Sessions' },
            ],
            'rows': [
            ],
        },
        'filter': {
        },
    }
    networks = Network.objects.all()
    for network in networks:
        sessions = PeeringSession.objects.filter(peer_netixlan__net=network)
        calculated = _render_alerts({
            'count': {
                'possible': sessions.count(),
                'provisioned': sessions.filter(provisioning_state=PeeringSession.PROV_COMPLETE).count(),
                'established': sessions.filter(operational_state=PeeringSession.OPER_ESTABLISHED).count(),
            },
        })
        row = { 'fields': [] }
        row['fields'].append({
            'display': network.name,
        })
        row['fields'].append({
            'display': network.asn,
        })
        row['fields'].append({
            'display': network.irr_as_set,
        })
        row['fields'].append({
            'display': network.looking_glass,
            'link': network.looking_glass,
        })
        row['fields'].append({
            'display': network.policy_general,
        })
        row['fields'].append({
            'display': calculated['count']['possible'],
            'alert': calculated['alert']['possible'],
        })
        row['fields'].append({
            'display': calculated['count']['provisioned'],
            'alert': calculated['alert']['provisioned'],
        })
        row['fields'].append({
            'display': calculated['count']['established'],
            'alert': calculated['alert']['established'],
        })
        context['table']['rows'].append(row)
    return HttpResponse(template.render(context, request))

def routers(request, rtr_id):
    context = {
        'view': {
            'name': 'routers',
            'url': reverse('prngmgr-routers'),
        }
    }
    if request.method == 'GET':
        if rtr_id:
            template = loader.get_template('prngmgr/form.html')
            if int(rtr_id) == 0:
                form = PeeringRouterForm()
                return HttpResponse(form)
            else:
                try:
                    router = PeeringRouter.objects.get(id=rtr_id)
                except:
                    return HttpResponseNotFound(rtr_id)
                form = PeeringRouterForm(instance=router)
            context['form'] = form
        else:
            template = loader.get_template('prngmgr/table.html')
            table = {
                'name': 'routers',
                'title': 'Peering Routers',
                'cols': [
                    { 'title': 'Hostname' },
                    { 'title': 'Peering Interfaces' },
                    { 'title': 'Possible Sessions' },
                    { 'title': 'Provisioned Sessions' },
                    { 'title': 'Established Sessions' },
                ],
                'rows': [
                ],
                'filter': {
                },
            }
            routers = PeeringRouter.objects.all()
            for router in routers:
                interfaces = PeeringRouterIXInterface.objects.filter(prngrtr=router)
                sessions = PeeringSession.objects.filter(prngrtriface__prngrtr=router)
                calculated = _render_alerts({
                    'count': {
                        'interfaces': interfaces.count(),
                        'possible': sessions.count(),
                        'provisioned': sessions.filter(provisioning_state=PeeringSession.PROV_COMPLETE).count(),
                        'established': sessions.filter(operational_state=PeeringSession.OPER_ESTABLISHED).count(),
                    },
                })
                row = { 'fields': [] }
                row['fields'].append({
                    'display': router.hostname,
                })
                row['fields'].append({
                    'display': calculated['count']['interfaces'],
                })
                row['fields'].append({
                    'display': calculated['count']['possible'],
                    'alert': calculated['alert']['possible'],
                })
                row['fields'].append({
                    'display': calculated['count']['provisioned'],
                    'alert': calculated['alert']['provisioned'],
                })
                row['fields'].append({
                    'display': calculated['count']['established'],
                    'alert': calculated['alert']['established'],
                })
                table['rows'].append(row)
            context['table'] = table
        return HttpResponse(template.render(context, request))

def ixps(request, ixp_id):
    template = loader.get_template('prngmgr/table.html')
    context = {
        'view': 'ixps',
        'table': {
            'name': 'ixps',
            'title': 'Internet Exchange Points',
            'cols': [
                { 'title': 'IXP Name' },
                { 'title': 'Country' },
                { 'title': 'Region' },
                { 'title': 'Participants' },
                { 'title': 'Possible Sessions' },
                { 'title': 'Provisioned Sessions' },
                { 'title': 'Established Sessions' },
            ],
            'rows': [
            ],
        },
        'filter': {
        },
    }
    ixps = InternetExchange.objects.all()
    for ixp in ixps:
        participants = NetworkIXLan.objects.filter(ixlan__ix=ixp).aggregate(count=Count('asn', distinct=True))
        sessions = PeeringSession.objects.filter(prngrtriface__netixlan__ixlan__ix=ixp)
        calculated = _render_alerts({
            'count': {
                'participants': participants['count'],
                'possible': sessions.count(),
                'provisioned': sessions.filter(provisioning_state=PeeringSession.PROV_COMPLETE).count(),
                'established': sessions.filter(operational_state=PeeringSession.OPER_ESTABLISHED).count(),
            },
        })
        row = { 'fields': [] }
        row['fields'].append({
            'display': ixp.name,
        })
        row['fields'].append({
            'display': ixp.country,
        })
        row['fields'].append({
            'display': ixp.region_continent,
        })
        row['fields'].append({
            'display': calculated['count']['participants'],
        })
        row['fields'].append({
            'display': calculated['count']['possible'],
            'alert': calculated['alert']['possible'],
        })
        row['fields'].append({
            'display': calculated['count']['provisioned'],
            'alert': calculated['alert']['provisioned'],
        })
        row['fields'].append({
            'display': calculated['count']['established'],
            'alert': calculated['alert']['established'],
        })
        context['table']['rows'].append(row)
    return HttpResponse(template.render(context, request))

def _render_alerts(calculated):
    if calculated['count']['possible'] == 0:
        calculated['alert'] = {
            'possible': ALERT_NONE,
            'provisioned': ALERT_NONE,
            'established': ALERT_NONE,
        }
    else:
        if calculated['count']['provisioned'] == 0:
            calculated['alert'] = {
                'possible': ALERT_SUCCESS,
                'provisioned': ALERT_DANGER,
                'established': ALERT_DANGER,
            }
        elif calculated['count']['provisioned'] < calculated['count']['possible']:
            if calculated['count']['established'] < calculated['count']['provisioned']:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_WARNING,
                    'established': ALERT_DANGER,
                }
            else:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_WARNING,
                    'established': ALERT_WARNING,
                }
        else:
            if calculated['count']['established'] < calculated['count']['provisioned']:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_SUCCESS,
                    'established': ALERT_DANGER,
                }
            else:
                calculated['alert'] = {
                    'possible': ALERT_SUCCESS,
                    'provisioned': ALERT_SUCCESS,
                    'established': ALERT_SUCCESS,
                }
    return calculated

<|endoftext|>


<commit_before>#!/usr/bin/python
#
# Copyright (c) 2015 Michael Scherer <mscherer@redhat.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# This script trigger a middleman rebuild if upstream git or any submodule
# did change

import yaml
import sys
import os
import subprocess
import datetime
import atexit
import syslog
import argparse

parser = argparse.ArgumentParser(description="Build middleman sites based "
                                             "on changes in git")
parser.add_argument("-f", "--force", help="force the rebuild",
                    action="store_true")
parser.add_argument("-n", "--dry-run", help="do not sync remotely",
                    action="store_true")
parser.add_argument("-d", "--debug", help="show debug output",
                    action="store_true")
parser.add_argument("config_file", help="yaml file for the builder config")

args = parser.parse_args()

def debug_print(message):
    if args.debug:
        print message

def refresh_checkout(checkout_dir):
    os.chdir(checkout_dir)
    subprocess.call(['git', 'fetch', '-q'])

def get_last_commit(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_last_commit_submodule(checkout_dir, submodule):
    os.chdir("{}/{}/".format(checkout_dir, submodule))
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_submodules_checkout(checkout_dir):
    os.chdir(checkout_dir)
    result = []
    submodule_status = subprocess.check_output(['git', 'submodule', 'status'])
    for s in submodule_status.split('\n'):
        # there is a empty line at the end...
        if s:
            result.append(s.split()[1])
    return result

def load_config(config_file):
    if not os.path.exists(config_file):
        print "Error %s, do not exist" % config_file
        sys.exit(1)
    if not os.path.isfile(config_file):
        print "Error %s is not a file" % config_file
        sys.exit(1)

    with open(config_file) as f:
        config = yaml.safe_load(f)

    return config

def has_submodules(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'submodule', 'status'])
    return len(r) > 0

# TODO complete that
def notify_error(stage, error):
    print error
    sys.exit(3)

if not args.config_file:
    print "This script take only 1 single argument, the config file"
    sys.exit(1)

config = load_config(args.config_file)

if not 'name' in config:
    print "Incorrect config file: {}".format(args.config_file)
    sys.exit(1)

name = config['name']

lock_file = os.path.expanduser('%s/lock_{}'.format(os.environ.get('XDG_RUNTIME_DIR', '~'), name))
if os.path.exists(lock_file):
    #TODO verify if the PID in the file still exist
    print "Builder already running, exiting"
    sys.exit(2)

# TODO try/except, show a better error message
fd = os.open(lock_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
os.write(fd, str(os.getpid()))
os.close(fd)
atexit.register(os.unlink, lock_file)

status_file = os.path.expanduser('~/status_%s.yml' % name)
status = {}
if os.path.exists(status_file):
    with open(status_file) as f:
        status = yaml.safe_load(f)

checkout_dir = os.path.expanduser("~/%s" % name)
if not os.path.isdir(checkout_dir):
    os.unlink(lock_file)
    print "Checkout not existing, exiting"
    sys.exit(2)

refresh_checkout(checkout_dir)

start_build = False

last_build = datetime.datetime.fromtimestamp(int(status.get("last_build", "0")))
if 'regular_rebuild_interval' in config and datetime.datetime.now() - last_build > datetime.timedelta(hours=config['regular_rebuild_interval']):
    start_build = True

current_commit = get_last_commit(checkout_dir)
if current_commit != status.get('last_build_commit', ''):
    start_build = True

submodule_commits = status.get('submodule_commits', {})
current_submodule_commits = {}
for submodule in get_submodules_checkout(checkout_dir):
    debug_print('Looking for %s' % submodule)
    current_submodule_commits[submodule] = get_last_commit_submodule(checkout_dir, submodule)

    if current_submodule_commits[submodule] != submodule_commits.get(submodule, ''):
        start_build = True

if args.force:
    start_build = True

if not start_build:
    debug_print("Nothing to build")
    sys.exit(0)

syslog.syslog("Start the build of {}".format(name))

os.chdir(checkout_dir)
if config.get('force_checkout', False):
    subprocess.call(['git', 'checkout', '-f'])
else:
    subprocess.call(['git', 'pull', '--rebase'])

if has_submodules(checkout_dir):
    subprocess.call(['git', 'submodule', 'init'])
    subprocess.call(['git', 'submodule', 'sync'])

if config.get('update_submodule_head', False):
    subprocess.call(['git', 'submodule', 'foreach',
                     '"git pull -qf origin master"'])

os.environ['PATH'] = "/usr/local/bin:/srv/builder/bin:" + os.environ['PATH']
try:
    syslog.syslog("Build of {}: bundle install".format(name))
    result = subprocess.check_output(['bundle', 'install'])
except subprocess.CalledProcessError, C:
    notify_error('install', C.output)

try:
    syslog.syslog("Build of {}: bundle exec middleman build".format(name))
    result = subprocess.check_output(['bundle', 'exec', 'middleman',
                                      'build', '--verbose'])
except subprocess.CalledProcessError, C:
    notify_error('build', C.output)

if not args.dry_run:
    syslog.syslog("Build of {}: start sync".format(name))
    # TODO log the message
    if config['remote']:
        subprocess.call(['rsync', '-e', 'ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i ' + os.path.expanduser('~/.ssh/{}_id.rsa'.format(name)), '--delete-after', '-rqavz', '%s/build/' % checkout_dir, config['remote']])
    else:
        subprocess.call(['bundle', 'exec', 'middleman', 'deploy'])
    syslog.syslog("Build of {}: finish sync".format(name))
else:
    syslog.syslog("Build of {}: not syncing, dry-run".format(name))

status = {}
status['last_build_commit'] = current_commit
status['last_build'] = datetime.datetime.now().strftime("%s")
status['submodule_commits'] = current_submodule_commits

with open(status_file, 'w+') as f:
    f.write(yaml.dump(status, default_flow_style=False))
<commit_msg>Fix bug and make pyflakes a bit more happy<commit_after>#!/usr/bin/python
#
# Copyright (c) 2015 Michael Scherer <mscherer@redhat.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# This script trigger a middleman rebuild if upstream git or any submodule
# did change

import yaml
import sys
import os
import subprocess
import datetime
import atexit
import syslog
import argparse

parser = argparse.ArgumentParser(description="Build middleman sites based "
                                             "on changes in git")
parser.add_argument("-f", "--force", help="force the rebuild",
                    action="store_true")
parser.add_argument("-n", "--dry-run", help="do not sync remotely",
                    action="store_true")
parser.add_argument("-d", "--debug", help="show debug output",
                    action="store_true")
parser.add_argument("config_file", help="yaml file for the builder config")

args = parser.parse_args()

def debug_print(message):
    if args.debug:
        print message

def refresh_checkout(checkout_dir):
    os.chdir(checkout_dir)
    subprocess.call(['git', 'fetch', '-q'])

def get_last_commit(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_last_commit_submodule(checkout_dir, submodule):
    os.chdir("{}/{}/".format(checkout_dir, submodule))
    r = subprocess.check_output(['git', 'ls-remote', '-q', '.',
                                 'refs/remotes/origin/HEAD'])
    return r.split()[0]

def get_submodules_checkout(checkout_dir):
    os.chdir(checkout_dir)
    result = []
    submodule_status = subprocess.check_output(['git', 'submodule', 'status'])
    for s in submodule_status.split('\n'):
        # there is a empty line at the end...
        if s:
            result.append(s.split()[1])
    return result

def load_config(config_file):
    if not os.path.exists(config_file):
        print "Error %s, do not exist" % config_file
        sys.exit(1)
    if not os.path.isfile(config_file):
        print "Error %s is not a file" % config_file
        sys.exit(1)

    with open(config_file) as f:
        config = yaml.safe_load(f)

    return config

def has_submodules(checkout_dir):
    os.chdir(checkout_dir)
    r = subprocess.check_output(['git', 'submodule', 'status'])
    return len(r) > 0

# TODO complete that
def notify_error(stage, error):
    print error
    sys.exit(3)

if not args.config_file:
    print "This script take only 1 single argument, the config file"
    sys.exit(1)

config = load_config(args.config_file)

if 'name' not in config:
    print "Incorrect config file: {}".format(args.config_file)
    sys.exit(1)

name = config['name']

lock_file = os.path.expanduser('{}/lock_{}'.format(
    os.environ.get('XDG_RUNTIME_DIR', '~'), name))
if os.path.exists(lock_file):
    # TODO verify if the PID in the file still exist
    print "Builder already running, exiting"
    sys.exit(2)

# TODO try/except, show a better error message
fd = os.open(lock_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
os.write(fd, str(os.getpid()))
os.close(fd)
atexit.register(os.unlink, lock_file)

status_file = os.path.expanduser('~/status_%s.yml' % name)
status = {}
if os.path.exists(status_file):
    with open(status_file) as f:
        status = yaml.safe_load(f)

checkout_dir = os.path.expanduser("~/%s" % name)
if not os.path.isdir(checkout_dir):
    os.unlink(lock_file)
    print "Checkout not existing, exiting"
    sys.exit(2)

refresh_checkout(checkout_dir)

start_build = False

last_build = datetime.datetime.fromtimestamp(
    int(status.get("last_build", "0")))

if ('regular_rebuild_interval' in config
        and datetime.datetime.now() - last_build
        > datetime.timedelta(hours=config['regular_rebuild_interval'])):
    start_build = True

current_commit = get_last_commit(checkout_dir)
if current_commit != status.get('last_build_commit', ''):
    start_build = True

submodule_commits = status.get('submodule_commits', {})
current_submodule_commits = {}
for submodule in get_submodules_checkout(checkout_dir):
    debug_print('Looking for %s' % submodule)
    current_submodule_commits[submodule] = \
        get_last_commit_submodule(checkout_dir, submodule)

    if current_submodule_commits[submodule] != submodule_commits.get(submodule, ''):
        start_build = True

if args.force:
    start_build = True

if not start_build:
    debug_print("Nothing to build")
    sys.exit(0)

syslog.syslog("Start the build of {}".format(name))

os.chdir(checkout_dir)
if config.get('force_checkout', False):
    subprocess.call(['git', 'checkout', '-f'])
else:
    subprocess.call(['git', 'pull', '--rebase'])

if has_submodules(checkout_dir):
    subprocess.call(['git', 'submodule', 'init'])
    subprocess.call(['git', 'submodule', 'sync'])

if config.get('update_submodule_head', False):
    subprocess.call(['git', 'submodule', 'foreach',
                     '"git pull -qf origin master"'])

os.environ['PATH'] = "/usr/local/bin:/srv/builder/bin:" + os.environ['PATH']
try:
    syslog.syslog("Build of {}: bundle install".format(name))
    result = subprocess.check_output(['bundle', 'install'])
except subprocess.CalledProcessError, C:
    notify_error('install', C.output)

try:
    syslog.syslog("Build of {}: bundle exec middleman build".format(name))
    result = subprocess.check_output(['bundle', 'exec', 'middleman',
                                      'build', '--verbose'])
except subprocess.CalledProcessError, C:
    notify_error('build', C.output)

if not args.dry_run:
    syslog.syslog("Build of {}: start sync".format(name))
    # TODO log the message
    if config['remote']:
        subprocess.call(['rsync', '-e', 'ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i ' + os.path.expanduser('~/.ssh/{}_id.rsa'.format(name)), '--delete-after', '-rqavz', '%s/build/' % checkout_dir, config['remote']])
    else:
        subprocess.call(['bundle', 'exec', 'middleman', 'deploy'])
    syslog.syslog("Build of {}: finish sync".format(name))
else:
    syslog.syslog("Build of {}: not syncing, dry-run".format(name))

status = {}
status['last_build_commit'] = current_commit
status['last_build'] = datetime.datetime.now().strftime("%s")
status['submodule_commits'] = current_submodule_commits

with open(status_file, 'w+') as f:
    f.write(yaml.dump(status, default_flow_style=False))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2015, 2016 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Because otherwise 'resource' collides with synapse.metrics.resource
from __future__ import absolute_import

import logging
from resource import getrusage, RUSAGE_SELF
import functools
import os
import stat
import time
import gc

from twisted.internet import reactor

from .metric import (
    CounterMetric, CallbackMetric, DistributionMetric, CacheMetric,
    MemoryUsageMetric,
)

logger = logging.getLogger(__name__)

all_metrics = []

class Metrics(object):
    """ A single Metrics object gives a (mutable) slice view of the all_metrics
    dict, allowing callers to easily register new metrics that are namespaced
    nicely."""

    def __init__(self, name):
        self.name_prefix = name

    def _register(self, metric_class, name, *args, **kwargs):
        full_name = "%s_%s" % (self.name_prefix, name)

        metric = metric_class(full_name, *args, **kwargs)

        all_metrics.append(metric)
        return metric

    def register_counter(self, *args, **kwargs):
        return self._register(CounterMetric, *args, **kwargs)

    def register_callback(self, *args, **kwargs):
        return self._register(CallbackMetric, *args, **kwargs)

    def register_distribution(self, *args, **kwargs):
        return self._register(DistributionMetric, *args, **kwargs)

    def register_cache(self, *args, **kwargs):
        return self._register(CacheMetric, *args, **kwargs)

def register_memory_metrics(hs):
    try:
        import psutil
        process = psutil.Process()
        process.memory_info().rss
    except (ImportError, AttributeError):
        logger.warn(
            "psutil is not installed or incorrect version."
            " Disabling memory metrics."
        )
        return
    metric = MemoryUsageMetric(hs, psutil)
    all_metrics.append(metric)

def get_metrics_for(pkg_name):
    """ Returns a Metrics instance for conveniently creating metrics
    namespaced with the given name prefix. """

    # Convert a "package.name" to "package_name" because Prometheus doesn't
    # let us use . in metric names
    return Metrics(pkg_name.replace(".", "_"))

def render_all():
    strs = []

    # TODO(paul): Internal hack
    update_resource_metrics()

    for metric in all_metrics:
        try:
            strs += metric.render()
        except Exception:
            strs += ["# FAILED to render"]
            logger.exception("Failed to render metric")

    strs.append("")  # to generate a final CRLF

    return "\n".join(strs)

# Now register some standard process-wide state metrics, to give indications of
# process resource usage

TICKS_PER_SEC = 100

rusage = None
stats = None

def update_resource_metrics():
    global rusage
    rusage = getrusage(RUSAGE_SELF)

    global stats
    with open("/proc/self/stat") as s:
        line = s.read()
        # line is PID (command) more stats go here ...
        stats = line.split(") ", 1)[1].split(" ")

## Legacy synapse-invented metric names

resource_metrics = get_metrics_for("process.resource")

# msecs
resource_metrics.register_callback("utime", lambda: rusage.ru_utime * 1000)
resource_metrics.register_callback("stime", lambda: rusage.ru_stime * 1000)

# kilobytes
resource_metrics.register_callback("maxrss", lambda: rusage.ru_maxrss * 1024)

TYPES = {
    stat.S_IFSOCK: "SOCK",
    stat.S_IFLNK: "LNK",
    stat.S_IFREG: "REG",
    stat.S_IFBLK: "BLK",
    stat.S_IFDIR: "DIR",
    stat.S_IFCHR: "CHR",
    stat.S_IFIFO: "FIFO",
}

def _process_fds():
    counts = {(k,): 0 for k in TYPES.values()}
    counts[("other",)] = 0

    # Not every OS will have a /proc/self/fd directory
    if not os.path.exists("/proc/self/fd"):
        return counts

    for fd in os.listdir("/proc/self/fd"):
        try:
            s = os.stat("/proc/self/fd/%s" % (fd))
            fmt = stat.S_IFMT(s.st_mode)
            if fmt in TYPES:
                t = TYPES[fmt]
            else:
                t = "other"

            counts[(t,)] += 1
        except OSError:
            # the dirh itself used by listdir() is usually missing by now
            pass

    return counts

get_metrics_for("process").register_callback("fds", _process_fds, labels=["type"])

## New prometheus-standard metric names
process_metrics = get_metrics_for("process");

process_metrics.register_callback(
    "cpu_user_seconds_total", lambda: float(stats[11]) / TICKS_PER_SEC
)
process_metrics.register_callback(
    "cpu_system_seconds_total", lambda: float(stats[12]) / TICKS_PER_SEC
)
process_metrics.register_callback(
    "cpu_seconds_total", lambda: (float(stats[11]) + float(stats[12])) / TICKS_PER_SEC
)

reactor_metrics = get_metrics_for("reactor")
tick_time = reactor_metrics.register_distribution("tick_time")
pending_calls_metric = reactor_metrics.register_distribution("pending_calls")

gc_time = reactor_metrics.register_distribution("gc_time", labels=["gen"])
gc_unreachable = reactor_metrics.register_counter("gc_unreachable", labels=["gen"])

reactor_metrics.register_callback(
    "gc_counts", lambda: {(i,): v for i, v in enumerate(gc.get_count())}, labels=["gen"]
)

def runUntilCurrentTimer(func):

    @functools.wraps(func)
    def f(*args, **kwargs):
        now = reactor.seconds()
        num_pending = 0

        # _newTimedCalls is one long list of *all* pending calls. Below loop
        # is based off of impl of reactor.runUntilCurrent
        for delayed_call in reactor._newTimedCalls:
            if delayed_call.time > now:
                break

            if delayed_call.delayed_time > 0:
                continue

            num_pending += 1

        num_pending += len(reactor.threadCallQueue)

        start = time.time() * 1000
        ret = func(*args, **kwargs)
        end = time.time() * 1000
        tick_time.inc_by(end - start)
        pending_calls_metric.inc_by(num_pending)

        # Check if we need to do a manual GC (since its been disabled), and do
        # one if necessary.
        threshold = gc.get_threshold()
        counts = gc.get_count()
        for i in (2, 1, 0):
            if threshold[i] < counts[i]:
                logger.info("Collecting gc %d", i)

                start = time.time() * 1000
                unreachable = gc.collect(i)
                end = time.time() * 1000

                gc_time.inc_by(end - start, i)
                gc_unreachable.inc_by(unreachable, i)

        return ret

    return f

try:
    # Ensure the reactor has all the attributes we expect
    reactor.runUntilCurrent
    reactor._newTimedCalls
    reactor.threadCallQueue

    # runUntilCurrent is called when we have pending calls. It is called once
    # per iteratation after fd polling.
    reactor.runUntilCurrent = runUntilCurrentTimer(reactor.runUntilCurrent)

    # We manually run the GC each reactor tick so that we can get some metrics
    # about time spent doing GC,
    gc.disable()
except AttributeError:
    pass
<commit_msg>Add standard process_*_memory_bytes metrics<commit_after># -*- coding: utf-8 -*-
# Copyright 2015, 2016 OpenMarket Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Because otherwise 'resource' collides with synapse.metrics.resource
from __future__ import absolute_import

import logging
from resource import getrusage, RUSAGE_SELF
import functools
import os
import stat
import time
import gc

from twisted.internet import reactor

from .metric import (
    CounterMetric, CallbackMetric, DistributionMetric, CacheMetric,
    MemoryUsageMetric,
)

logger = logging.getLogger(__name__)

all_metrics = []

class Metrics(object):
    """ A single Metrics object gives a (mutable) slice view of the all_metrics
    dict, allowing callers to easily register new metrics that are namespaced
    nicely."""

    def __init__(self, name):
        self.name_prefix = name

    def _register(self, metric_class, name, *args, **kwargs):
        full_name = "%s_%s" % (self.name_prefix, name)

        metric = metric_class(full_name, *args, **kwargs)

        all_metrics.append(metric)
        return metric

    def register_counter(self, *args, **kwargs):
        return self._register(CounterMetric, *args, **kwargs)

    def register_callback(self, *args, **kwargs):
        return self._register(CallbackMetric, *args, **kwargs)

    def register_distribution(self, *args, **kwargs):
        return self._register(DistributionMetric, *args, **kwargs)

    def register_cache(self, *args, **kwargs):
        return self._register(CacheMetric, *args, **kwargs)

def register_memory_metrics(hs):
    try:
        import psutil
        process = psutil.Process()
        process.memory_info().rss
    except (ImportError, AttributeError):
        logger.warn(
            "psutil is not installed or incorrect version."
            " Disabling memory metrics."
        )
        return
    metric = MemoryUsageMetric(hs, psutil)
    all_metrics.append(metric)

def get_metrics_for(pkg_name):
    """ Returns a Metrics instance for conveniently creating metrics
    namespaced with the given name prefix. """

    # Convert a "package.name" to "package_name" because Prometheus doesn't
    # let us use . in metric names
    return Metrics(pkg_name.replace(".", "_"))

def render_all():
    strs = []

    # TODO(paul): Internal hack
    update_resource_metrics()

    for metric in all_metrics:
        try:
            strs += metric.render()
        except Exception:
            strs += ["# FAILED to render"]
            logger.exception("Failed to render metric")

    strs.append("")  # to generate a final CRLF

    return "\n".join(strs)

# Now register some standard process-wide state metrics, to give indications of
# process resource usage

TICKS_PER_SEC = 100
BYTES_PER_PAGE = 4096

rusage = None
stats = None

def update_resource_metrics():
    global rusage
    rusage = getrusage(RUSAGE_SELF)

    global stats
    with open("/proc/self/stat") as s:
        line = s.read()
        # line is PID (command) more stats go here ...
        stats = line.split(") ", 1)[1].split(" ")

## Legacy synapse-invented metric names

resource_metrics = get_metrics_for("process.resource")

# msecs
resource_metrics.register_callback("utime", lambda: rusage.ru_utime * 1000)
resource_metrics.register_callback("stime", lambda: rusage.ru_stime * 1000)

# kilobytes
resource_metrics.register_callback("maxrss", lambda: rusage.ru_maxrss * 1024)

TYPES = {
    stat.S_IFSOCK: "SOCK",
    stat.S_IFLNK: "LNK",
    stat.S_IFREG: "REG",
    stat.S_IFBLK: "BLK",
    stat.S_IFDIR: "DIR",
    stat.S_IFCHR: "CHR",
    stat.S_IFIFO: "FIFO",
}

def _process_fds():
    counts = {(k,): 0 for k in TYPES.values()}
    counts[("other",)] = 0

    # Not every OS will have a /proc/self/fd directory
    if not os.path.exists("/proc/self/fd"):
        return counts

    for fd in os.listdir("/proc/self/fd"):
        try:
            s = os.stat("/proc/self/fd/%s" % (fd))
            fmt = stat.S_IFMT(s.st_mode)
            if fmt in TYPES:
                t = TYPES[fmt]
            else:
                t = "other"

            counts[(t,)] += 1
        except OSError:
            # the dirh itself used by listdir() is usually missing by now
            pass

    return counts

get_metrics_for("process").register_callback("fds", _process_fds, labels=["type"])

## New prometheus-standard metric names
process_metrics = get_metrics_for("process");

process_metrics.register_callback(
    "cpu_user_seconds_total", lambda: float(stats[11]) / TICKS_PER_SEC
)
process_metrics.register_callback(
    "cpu_system_seconds_total", lambda: float(stats[12]) / TICKS_PER_SEC
)
process_metrics.register_callback(
    "cpu_seconds_total", lambda: (float(stats[11]) + float(stats[12])) / TICKS_PER_SEC
)

process_metrics.register_callback(
    "virtual_memory_bytes", lambda: int(stats[20])
)
process_metrics.register_callback(
    "resident_memory_bytes", lambda: int(stats[21]) * BYTES_PER_PAGE
)

reactor_metrics = get_metrics_for("reactor")
tick_time = reactor_metrics.register_distribution("tick_time")
pending_calls_metric = reactor_metrics.register_distribution("pending_calls")

gc_time = reactor_metrics.register_distribution("gc_time", labels=["gen"])
gc_unreachable = reactor_metrics.register_counter("gc_unreachable", labels=["gen"])

reactor_metrics.register_callback(
    "gc_counts", lambda: {(i,): v for i, v in enumerate(gc.get_count())}, labels=["gen"]
)

def runUntilCurrentTimer(func):

    @functools.wraps(func)
    def f(*args, **kwargs):
        now = reactor.seconds()
        num_pending = 0

        # _newTimedCalls is one long list of *all* pending calls. Below loop
        # is based off of impl of reactor.runUntilCurrent
        for delayed_call in reactor._newTimedCalls:
            if delayed_call.time > now:
                break

            if delayed_call.delayed_time > 0:
                continue

            num_pending += 1

        num_pending += len(reactor.threadCallQueue)

        start = time.time() * 1000
        ret = func(*args, **kwargs)
        end = time.time() * 1000
        tick_time.inc_by(end - start)
        pending_calls_metric.inc_by(num_pending)

        # Check if we need to do a manual GC (since its been disabled), and do
        # one if necessary.
        threshold = gc.get_threshold()
        counts = gc.get_count()
        for i in (2, 1, 0):
            if threshold[i] < counts[i]:
                logger.info("Collecting gc %d", i)

                start = time.time() * 1000
                unreachable = gc.collect(i)
                end = time.time() * 1000

                gc_time.inc_by(end - start, i)
                gc_unreachable.inc_by(unreachable, i)

        return ret

    return f

try:
    # Ensure the reactor has all the attributes we expect
    reactor.runUntilCurrent
    reactor._newTimedCalls
    reactor.threadCallQueue

    # runUntilCurrent is called when we have pending calls. It is called once
    # per iteratation after fd polling.
    reactor.runUntilCurrent = runUntilCurrentTimer(reactor.runUntilCurrent)

    # We manually run the GC each reactor tick so that we can get some metrics
    # about time spent doing GC,
    gc.disable()
except AttributeError:
    pass
<|endoftext|>