path: root/library/cpp/messagebus/locator.cpp

////////////////////////////////////////////////////////////////////////////
/// \file 
/// \brief Implementation of locator service

#include "locator.h"

#include "ybus.h"

#include <util/generic/hash_set.h>
#include <util/system/hostname.h>

namespace NBus {
    using namespace NAddr; 

    static TIpPort GetAddrPort(const IRemoteAddr& addr) { 
        switch (addr.Addr()->sa_family) { 
            case AF_INET: { 
                return ntohs(((const sockaddr_in*)addr.Addr())->sin_port); 
            } 

            case AF_INET6: { 
                return ntohs(((const sockaddr_in6*)addr.Addr())->sin6_port); 
            } 
 
            default: { 
                ythrow yexception() << "not implemented"; 
                break; 
            } 
        }
    } 

    static inline bool GetIp6AddressFromVector(const TVector<TNetAddr>& addrs, TNetAddr* addr) { 
        for (size_t i = 1; i < addrs.size(); ++i) { 
            if (addrs[i - 1].Addr()->sa_family == addrs[i].Addr()->sa_family) { 
                return false; 
            } 
 
            if (GetAddrPort(addrs[i - 1]) != GetAddrPort(addrs[i])) { 
                return false; 
            } 
        }

        for (size_t i = 0; i < addrs.size(); ++i) { 
            if (addrs[i].Addr()->sa_family == AF_INET6) { 
                *addr = addrs[i]; 
                return true; 
            } 
        }
 
        return false; 
    }

    EMessageStatus TBusProtocol::GetDestination(const TBusClientSession*, TBusMessage* mess, TBusLocator* locator, TNetAddr* addr) { 
        TBusService service = GetService(); 
        TBusKey key = GetKey(mess); 
        TVector<TNetAddr> addrs; 
 
        /// check for special local key 
        if (key == YBUS_KEYLOCAL) { 
            locator->GetLocalAddresses(service, addrs); 
        } else { 
            /// lookup address/port in the locator table 
            locator->LocateAll(service, key, addrs); 
        }

        if (addrs.size() == 0) { 
            return MESSAGE_SERVICE_UNKNOWN; 
        } else if (addrs.size() == 1) { 
            *addr = addrs[0]; 
        } else { 
            if (!GetIp6AddressFromVector(addrs, addr)) { 
                /// default policy can't make choice for you here, overide GetDestination() function 
                /// to implement custom routing strategy for your service. 
                return MESSAGE_SERVICE_TOOMANY; 
            } 
        }
 
        return MESSAGE_OK; 
    }

    static const sockaddr_in* SockAddrIpV4(const IRemoteAddr& a) { 
        return (const sockaddr_in*)a.Addr(); 
    }

    static const sockaddr_in6* SockAddrIpV6(const IRemoteAddr& a) { 
        return (const sockaddr_in6*)a.Addr(); 
    }

    static bool IsAddressEqual(const IRemoteAddr& a1, const IRemoteAddr& a2) { 
        if (a1.Addr()->sa_family == a2.Addr()->sa_family) { 
            if (a1.Addr()->sa_family == AF_INET) { 
                return memcmp(&SockAddrIpV4(a1)->sin_addr, &SockAddrIpV4(a2)->sin_addr, sizeof(in_addr)) == 0; 
            } else { 
                return memcmp(&SockAddrIpV6(a1)->sin6_addr, &SockAddrIpV6(a2)->sin6_addr, sizeof(in6_addr)) == 0; 
            } 
        }
        return false; 
    }

    TBusLocator::TBusLocator() 
        : MyInterfaces(GetNetworkInterfaces()) 
    { 
    } 

    bool TBusLocator::TItem::operator<(const TItem& y) const { 
        const TItem& x = *this; 

        if (x.ServiceId == y.ServiceId) { 
            return (x.End < y.End) || ((x.End == y.End) && CompareByHost(x.Addr, y.Addr) < 0); 
        }
        return x.ServiceId < y.ServiceId; 
    }

    bool TBusLocator::TItem::operator==(const TItem& y) const { 
        return ServiceId == y.ServiceId && Start == y.Start && End == y.End && Addr == y.Addr; 
    } 

    TBusLocator::TItem::TItem(TServiceId serviceId, TBusKey start, TBusKey end, const TNetAddr& addr) 
        : ServiceId(serviceId) 
        , Start(start) 
        , End(end) 
        , Addr(addr) 
    { 
    }

    bool TBusLocator::IsLocal(const TNetAddr& addr) { 
        for (const auto& myInterface : MyInterfaces) { 
            if (IsAddressEqual(addr, *myInterface.Address)) { 
                return true; 
            } 
        } 

        return false; 
    } 

    TBusLocator::TServiceId TBusLocator::GetServiceId(const char* name) { 
        const char* c = ServiceIdSet.insert(name).first->c_str(); 
        return (ui64)c; 
    }

    int TBusLocator::RegisterBreak(TBusService service, const TVector<TBusKey>& starts, const TNetAddr& addr) { 
        TGuard<TMutex> G(Lock); 

        TServiceId serviceId = GetServiceId(service); 
        for (size_t i = 0; i < starts.size(); ++i) { 
            RegisterBreak(serviceId, starts[i], addr); 
        } 
        return 0; 
    }

    int TBusLocator::RegisterBreak(TServiceId serviceId, const TBusKey start, const TNetAddr& addr) { 
        TItems::const_iterator it = Items.lower_bound(TItem(serviceId, 0, start, addr)); 
        TItems::const_iterator service_it = 
            Items.lower_bound(TItem(serviceId, 0, 0, TNetAddr()));

        THolder<TItem> left; 
        THolder<TItem> right; 
        if ((it != Items.end() || Items.begin() != Items.end()) && service_it != Items.end() && service_it->ServiceId == serviceId) { 
            if (it == Items.end()) { 
                --it; 
            } 
            const TItem& item = *it; 
            left.Reset(new TItem(serviceId, item.Start, 
                                 Max<TBusKey>(item.Start, start - 1), item.Addr)); 
            right.Reset(new TItem(serviceId, start, item.End, addr)); 
            Items.erase(*it); 
        } else { 
            left.Reset(new TItem(serviceId, YBUS_KEYMIN, start, addr)); 
            if (start < YBUS_KEYMAX) { 
                right.Reset(new TItem(serviceId, start + 1, YBUS_KEYMAX, addr)); 
            } 
        }
        Items.insert(*left); 
        Items.insert(*right); 
        NormalizeBreaks(serviceId); 
        return 0; 
    }

    int TBusLocator::UnregisterBreak(TBusService service, const TNetAddr& addr) { 
        TGuard<TMutex> G(Lock); 

        TServiceId serviceId = GetServiceId(service); 
        return UnregisterBreak(serviceId, addr); 
    } 

    int TBusLocator::UnregisterBreak(TServiceId serviceId, const TNetAddr& addr) { 
        int deleted = 0; 
        TItems::iterator it = Items.begin(); 
        while (it != Items.end()) { 
            const TItem& item = *it; 
            if (item.ServiceId != serviceId) { 
                ++it; 
                continue; 
            } 
            TItems::iterator itErase = it++; 
            if (item.ServiceId == serviceId && item.Addr == addr) { 
                Items.erase(itErase); 
                deleted += 1; 
            } 
        } 

        if (Items.begin() == Items.end()) { 
            return deleted; 
        } 
        TBusKey keyItem = YBUS_KEYMAX; 
        it = Items.end(); 
        TItems::iterator first = it; 
        do { 
            --it; 
            // item.Start is not used in set comparison function 
            // so you can't violate set sort order by changing it 
            // hence const_cast() 
            TItem& item = const_cast<TItem&>(*it); 
            if (item.ServiceId != serviceId) { 
                continue; 
            } 
            first = it; 
            if (item.End < keyItem) { 
                item.End = keyItem; 
            } 
            keyItem = item.Start - 1; 
        } while (it != Items.begin()); 

        if (first != Items.end() && first->Start != 0) { 
            TItem item(serviceId, YBUS_KEYMIN, first->Start - 1, first->Addr); 
            Items.insert(item); 
        }

        NormalizeBreaks(serviceId); 
        return deleted;
    }

    void TBusLocator::NormalizeBreaks(TServiceId serviceId) { 
        TItems::const_iterator first = Items.lower_bound(TItem(serviceId, YBUS_KEYMIN, YBUS_KEYMIN, TNetAddr())); 
        TItems::const_iterator last = Items.end(); 

        if ((Items.end() != first) && (first->ServiceId == serviceId)) { 
            if (serviceId != Max<TServiceId>()) { 
                last = Items.lower_bound(TItem(serviceId + 1, YBUS_KEYMIN, YBUS_KEYMIN, TNetAddr())); 
            } 

            --last; 
            Y_ASSERT(Items.end() != last); 
            Y_ASSERT(last->ServiceId == serviceId); 

            TItem& beg = const_cast<TItem&>(*first); 
            beg.Addr = last->Addr; 
        }
    }

    int TBusLocator::LocateAll(TBusService service, TBusKey key, TVector<TNetAddr>& addrs) { 
        TGuard<TMutex> G(Lock); 
        Y_VERIFY(addrs.empty(), "Non emtpy addresses"); 

        TServiceId serviceId = GetServiceId(service); 
        TItems::const_iterator it; 

        for (it = Items.lower_bound(TItem(serviceId, 0, key, TNetAddr())); 
             it != Items.end() && it->ServiceId == serviceId && it->Start <= key && key <= it->End; 
             ++it) { 
            const TItem& item = *it; 
            addrs.push_back(item.Addr); 
        } 

        if (addrs.size() == 0) { 
            return -1; 
        } 
        return (int)addrs.size(); 
    }
 
    int TBusLocator::Locate(TBusService service, TBusKey key, TNetAddr* addr) { 
        TGuard<TMutex> G(Lock); 

        TServiceId serviceId = GetServiceId(service); 
        TItems::const_iterator it; 

        it = Items.lower_bound(TItem(serviceId, 0, key, TNetAddr())); 

        if (it != Items.end()) { 
            const TItem& item = *it; 
            if (item.ServiceId == serviceId && item.Start <= key && key < item.End) { 
                *addr = item.Addr; 

                return 0; 
            } 
        } 
 
        return -1; 
    }

    int TBusLocator::GetLocalPort(TBusService service) { 
        TGuard<TMutex> G(Lock); 
        TServiceId serviceId = GetServiceId(service); 
        TItems::const_iterator it; 
        int port = 0; 

        for (it = Items.lower_bound(TItem(serviceId, 0, 0, TNetAddr())); it != Items.end(); ++it) { 
            const TItem& item = *it; 
            if (item.ServiceId != serviceId) { 
                break; 
            } 

            if (IsLocal(item.Addr)) { 
                if (port != 0 && port != GetAddrPort(item.Addr)) { 
                    Y_ASSERT(0 && "Can't decide which port to use."); 
                    return 0; 
                } 
                port = GetAddrPort(item.Addr); 
            } 
        }

        return port; 
    }

    int TBusLocator::GetLocalAddresses(TBusService service, TVector<TNetAddr>& addrs) { 
        TGuard<TMutex> G(Lock); 
        TServiceId serviceId = GetServiceId(service); 
        TItems::const_iterator it; 

        for (it = Items.lower_bound(TItem(serviceId, 0, 0, TNetAddr())); it != Items.end(); ++it) { 
            const TItem& item = *it; 
            if (item.ServiceId != serviceId) { 
                break; 
            } 

            if (IsLocal(item.Addr)) { 
                addrs.push_back(item.Addr); 
            } 
        }

        if (addrs.size() == 0) { 
            return -1; 
        }

        return (int)addrs.size(); 
    }

    int TBusLocator::LocateHost(TBusService service, TBusKey key, TString* host, int* port, bool* isLocal) { 
        int ret; 
        TNetAddr addr; 
        ret = Locate(service, key, &addr); 
        if (ret != 0) { 
            return ret; 
        } 

        { 
            TGuard<TMutex> G(Lock); 
            THostAddrMap::const_iterator it = HostAddrMap.find(addr); 
            if (it == HostAddrMap.end()) { 
                return -1; 
            } 
            *host = it->second; 
        } 
 
        *port = GetAddrPort(addr); 
        if (isLocal != nullptr) { 
            *isLocal = IsLocal(addr); 
        } 
        return 0; 
    }

    int TBusLocator::LocateKeys(TBusService service, TBusKeyVec& keys, bool onlyLocal) { 
        TGuard<TMutex> G(Lock);
        Y_VERIFY(keys.empty(), "Non empty keys"); 
 
        TServiceId serviceId = GetServiceId(service); 
        TItems::const_iterator it; 
        for (it = Items.begin(); it != Items.end(); ++it) { 
            const TItem& item = *it; 
            if (item.ServiceId != serviceId) { 
                continue; 
            } 
            if (onlyLocal && !IsLocal(item.Addr)) { 
                continue; 
            } 
            keys.push_back(std::make_pair(item.Start, item.End)); 
        }
        return (int)keys.size(); 
    }

    int TBusLocator::Register(TBusService service, const char* hostName, int port, TBusKey start /*= YBUS_KEYMIN*/, TBusKey end /*= YBUS_KEYMAX*/, EIpVersion requireVersion /*= EIP_VERSION_4*/, EIpVersion preferVersion /*= EIP_VERSION_ANY*/) { 
        TNetAddr addr(hostName, port, requireVersion, preferVersion); // throws 
        { 
            TGuard<TMutex> G(Lock); 
            HostAddrMap[addr] = hostName; 
        } 
        Register(service, start, end, addr); 
        return 0; 
    }

    int TBusLocator::Register(TBusService service, TBusKey start, TBusKey end, const TNetworkAddress& na, EIpVersion requireVersion /*= EIP_VERSION_4*/, EIpVersion preferVersion /*= EIP_VERSION_ANY*/) { 
        TNetAddr addr(na, requireVersion, preferVersion); // throws 
        Register(service, start, end, addr); 
        return 0; 
    }

    int TBusLocator::Register(TBusService service, TBusKey start, TBusKey end, const TNetAddr& addr) { 
        TGuard<TMutex> G(Lock);

        TServiceId serviceId = GetServiceId(service); 
        TItems::const_iterator it; 

        TItem itemToReg(serviceId, start, end, addr); 
        for (it = Items.lower_bound(TItem(serviceId, 0, start, TNetAddr())); 
             it != Items.end() && it->ServiceId == serviceId; 
             ++it) { 
            const TItem& item = *it; 
            if (item == itemToReg) { 
                return 0; 
            } 
            if ((item.Start < start && start < item.End) || (item.Start < end && end < item.End)) { 
                Y_FAIL("Overlap in registered keys with non-identical range"); 
            } 
        } 

        Items.insert(itemToReg); 
        return 0; 
    } 

    int TBusLocator::Unregister(TBusService service, TBusKey start, TBusKey end) { 
        TGuard<TMutex> G(Lock); 
        TServiceId serviceId = GetServiceId(service); 
        Items.erase(TItem(serviceId, start, end, TNetAddr())); 
        return 0; 
    }

}