// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.


#include <gflags/gflags.h>
#include "brpc/reloadable_flags.h"
#include "brpc/load_balancer.h"
#include "brpc/socket.h"


namespace brpc {

DEFINE_bool(show_lb_in_vars, false, "Describe LoadBalancers in vars");
DEFINE_int32(default_weight_of_wlb, 0, "Default weight value of Weighted LoadBalancer(wlb). "
             "wlb policy degradation is enabled when default_weight_of_wlb > 0 to avoid some "
             "problems when user is using wlb but forgot to set the weights of some of their "
             "downstream instances. Then these instances will be set default_weight_of_wlb as "
             "their weights. wlb policy degradation is not enabled by default.");
BRPC_VALIDATE_GFLAG(show_lb_in_vars, PassValidate);

// For assigning unique names for lb.
static butil::static_atomic<int> g_lb_counter = BUTIL_STATIC_ATOMIC_INIT(0);

bool LoadBalancer::IsServerAvailable(SocketId id, SocketUniquePtr* out) {
    SocketUniquePtr ptr;
    bool res = Socket::Address(id, &ptr) == 0 && ptr->IsAvailable();
    if (res) {
        *out = std::move(ptr);
    }
    return res;
}

void SharedLoadBalancer::DescribeLB(std::ostream& os, void* arg) {
    (static_cast<SharedLoadBalancer*>(arg))->Describe(os, DescribeOptions());
}

void SharedLoadBalancer::ExposeLB() {
    bool changed = false;
    _st_mutex.lock();
    if (!_exposed) {
        _exposed = true;
        changed = true;
    }
    _st_mutex.unlock();
    if (changed) {
        char name[32];
        snprintf(name, sizeof(name), "_load_balancer_%d", g_lb_counter.fetch_add(
                     1, butil::memory_order_relaxed));
        _st.expose(name);
    }
}

SharedLoadBalancer::SharedLoadBalancer()
    : _lb(NULL)
    , _weight_sum(0)
    , _exposed(false)
    , _st(DescribeLB, this) {
}

SharedLoadBalancer::~SharedLoadBalancer() {
    _st.hide();
    if (_lb) {
        _lb->Destroy();
        _lb = NULL;
    }
}

int SharedLoadBalancer::Init(const char* lb_protocol) {
    std::string lb_name;
    butil::StringPiece lb_params;
    if (!ParseParameters(lb_protocol, &lb_name, &lb_params)) {
        LOG(FATAL) << "Fail to parse this load balancer protocol '" << lb_protocol << '\'';
        return -1;
    }
    const LoadBalancer* lb = LoadBalancerExtension()->Find(lb_name.c_str());
    if (lb == NULL) {
        LOG(FATAL) << "Fail to find LoadBalancer by `" << lb_name << "'";
        return -1;
    }
    _lb = lb->New(lb_params);
    if (_lb == NULL) {
        LOG(FATAL) << "Fail to new LoadBalancer";
        return -1;
    }
    if (FLAGS_show_lb_in_vars && !_exposed) {
        ExposeLB();
    }
    return 0;
}

void SharedLoadBalancer::Describe(std::ostream& os,
                                  const DescribeOptions& options) {
    if (_lb == NULL) {
        os << "lb=NULL";
    } else {
        _lb->Describe(os, options);
    }
}

bool SharedLoadBalancer::ParseParameters(const butil::StringPiece& lb_protocol,
                                         std::string* lb_name,
                                         butil::StringPiece* lb_params) {
    lb_name->clear();
    lb_params->clear();
    if (lb_protocol.empty()) {
        return false;
    }
    const char separator = ':';
    size_t pos = lb_protocol.find(separator);
    if (pos == std::string::npos) {
        lb_name->append(lb_protocol.data(), lb_protocol.size());
    } else {
        lb_name->append(lb_protocol.data(), pos);
        if (pos < lb_protocol.size() - sizeof(separator)) {
            *lb_params = lb_protocol.substr(pos + sizeof(separator));
        }
    }

    return true;
}
																				 
} // namespace brpc
