installers/charm/zookeeper/src/charm.py

#!/usr/bin/env python3
# Copyright 2021 Canonical 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.
#
# For those usages not covered by the Apache License, Version 2.0 please
# contact: legal@canonical.com
#
# To get in touch with the maintainers, please contact:
# osm-charmers@lists.launchpad.net
##

# pylint: disable=E0213

import logging
from typing import NoReturn


from ops.framework import EventBase
from ops.main import main
from opslib.osm.charm import CharmedOsmBase
from opslib.osm.interfaces.zookeeper import ZookeeperCluster, ZookeeperServer
from opslib.osm.pod import ContainerV3Builder, PodSpecV3Builder
from opslib.osm.validator import ModelValidator, validator

logger = logging.getLogger(__name__)

CLIENT_PORT = 2181
SERVER_PORT = 2888
LEADER_ELECTION_PORT = 3888


class ConfigModel(ModelValidator):
    log_level: str
    image_pull_policy: str
    min_session_timeout: int
    max_session_timeout: int
    purge_interval: int
    snap_retain_count: int
    max_client_cnxns: int
    heap: int
    sync_limit: int
    init_limit: int
    tick_time: int
    security_context: bool

    @validator("log_level")
    def validate_log_level(cls, v):
        if v not in {"INFO", "DEBUG"}:
            raise ValueError("value must be INFO or DEBUG")
        return v

    @validator("image_pull_policy")
    def validate_image_pull_policy(cls, v):
        values = {
            "always": "Always",
            "ifnotpresent": "IfNotPresent",
            "never": "Never",
        }
        v = v.lower()
        if v not in values.keys():
            raise ValueError("value must be always, ifnotpresent or never")
        return values[v]


class ZookeeperCharm(CharmedOsmBase):
    """Zookeeper Charm."""

    def __init__(self, *args) -> NoReturn:
        """Zookeeper Charm constructor."""
        super().__init__(*args, oci_image="image")
        # Initialize Zookeeper cluster relation
        self.zookeeper_cluster = ZookeeperCluster(self, "cluster", CLIENT_PORT)
        self.framework.observe(self.on["cluster"].relation_changed, self._setup_cluster)
        # Initialize Zookeeper relation
        self.zookeeper_server = ZookeeperServer(self, "zookeeper")
        self.framework.observe(self.on["zookeeper"].relation_joined, self._publish_info)

    @property
    def num_units(self):
        return self.zookeeper_cluster.num_units

    @property
    def zookeeper_uri(self):
        return self.zookeeper_cluster.zookeeper_uri

    def _setup_cluster(self, event: EventBase):
        """Publishes Zookeeper information and reconfigures the pod.

        Args:
            event (EventBase): Zookeeper Cluster relation event.
        """
        self._publish_info(event)
        self.configure_pod()

    def _publish_info(self, event: EventBase):
        """Publishes Zookeeper information.

        Args:
            event (EventBase): Zookeeper relation event.
        """
        if self.unit.is_leader():
            zk_uri = self.zookeeper_uri
            if zk_uri:
                self.zookeeper_server.publish_info(zk_uri)
            else:
                event.defer()

    def build_pod_spec(self, image_info):
        # Validate config
        config = ConfigModel(**dict(self.config))

        # Create Builder for the PodSpec
        pod_spec_builder = PodSpecV3Builder(
            enable_security_context=config.security_context
        )

        # Build Container
        container_builder = ContainerV3Builder(
            self.app.name,
            image_info,
            config.image_pull_policy,
            run_as_non_root=config.security_context,
        )

        container_builder.add_port(name="client", port=CLIENT_PORT)
        container_builder.add_port(name="server", port=SERVER_PORT)
        container_builder.add_port(name="leader-election", port=LEADER_ELECTION_PORT)
        container_builder.add_tcpsocket_readiness_probe(
            CLIENT_PORT,
            initial_delay_seconds=10,
            timeout_seconds=5,
            failure_threshold=6,
            success_threshold=1,
        )
        container_builder.add_tcpsocket_liveness_probe(
            CLIENT_PORT, initial_delay_seconds=20
        )
        container_builder.add_command(
            [
                "sh",
                "-c",
                " ".join(
                    [
                        "start-zookeeper",
                        f"--servers={self.num_units}",
                        "--data_dir=/var/lib/zookeeper/data",
                        "--data_log_dir=/var/lib/zookeeper/data/log",
                        "--conf_dir=/opt/zookeeper/conf",
                        f"--client_port={CLIENT_PORT}",
                        f"--election_port={LEADER_ELECTION_PORT}",
                        f"--server_port={SERVER_PORT}",
                        f"--tick_time={config.tick_time}",
                        f"--init_limit={config.init_limit}",
                        f"--sync_limit={config.sync_limit}",
                        f"--heap={config.heap}M",
                        f"--max_client_cnxns={config.max_client_cnxns}",
                        f"--snap_retain_count={config.snap_retain_count}",
                        f"--purge_interval={config.purge_interval}",
                        f"--max_session_timeout={config.max_session_timeout}",
                        f"--min_session_timeout={config.min_session_timeout}",
                        f"--log_level={config.log_level}",
                    ]
                ),
            ]
        )

        container = container_builder.build()

        # Add container to pod spec
        pod_spec_builder.add_container(container)

        return pod_spec_builder.build()


if __name__ == "__main__":
    main(ZookeeperCharm)